fix: pass enable_network_isolation when creating TF and SKLearn models (#1043)

Author: laurenyu

src/sagemaker/sklearn/estimator.py

@@ -172,6 +172,7 @@ def create_model(
             image=self.image_name,
             sagemaker_session=self.sagemaker_session,
             vpc_config=self.get_vpc_config(vpc_config_override),
+            enable_network_isolation=self.enable_network_isolation(),
             **kwargs
         )
 

src/sagemaker/tensorflow/estimator.py

@@ -583,6 +583,7 @@ def _create_tfs_model(
             entry_point=entry_point,
             source_dir=source_dir,
             dependencies=dependencies,
+            enable_network_isolation=self.enable_network_isolation(),
         )
 
     def _create_default_model(
@@ -612,6 +613,7 @@ def _create_default_model(
             sagemaker_session=self.sagemaker_session,
             vpc_config=self.get_vpc_config(vpc_config_override),
             dependencies=dependencies or self.dependencies,
+            enable_network_isolation=self.enable_network_isolation(),
         )
 
     def hyperparameters(self):

tests/unit/test_sklearn.py

@@ -1,4 +1,4 @@
-# Copyright 2017-2018 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+# Copyright 2017-2019 Amazon.com, Inc. or its affiliates. All Rights Reserved.
 #
 # Licensed under the Apache License, Version 2.0 (the "License"). You
 # may not use this file except in compliance with the License. A copy of
@@ -20,12 +20,10 @@
 from mock import Mock
 from mock import patch
 
-
 from sagemaker.sklearn import defaults
 from sagemaker.sklearn import SKLearn
 from sagemaker.sklearn import SKLearnPredictor, SKLearnModel
 
-
 DATA_DIR = os.path.join(os.path.dirname(__file__), "..", "data")
 SCRIPT_PATH = os.path.join(DATA_DIR, "dummy_script.py")
 TIMESTAMP = "2017-11-06-14:14:15.672"
@@ -183,6 +181,7 @@ def test_create_model_from_estimator(sagemaker_session, sklearn_version):
         py_version=PYTHON_VERSION,
         base_job_name="job",
         source_dir=source_dir,
+        enable_network_isolation=True,
     )
 
     job_name = "new_name"
@@ -198,6 +197,7 @@ def test_create_model_from_estimator(sagemaker_session, sklearn_version):
     assert model.container_log_level == container_log_level
     assert model.source_dir == source_dir
     assert model.vpc_config is None
+    assert model.enable_network_isolation()
 
 
 def test_create_model_with_optional_params(sagemaker_session):

tests/unit/test_tf_estimator.py

@@ -22,7 +22,7 @@
 from sagemaker.fw_utils import create_image_uri
 from sagemaker.model import MODEL_SERVER_WORKERS_PARAM_NAME
 from sagemaker.session import s3_input
-from sagemaker.tensorflow import defaults, TensorFlow, TensorFlowModel, TensorFlowPredictor
+from sagemaker.tensorflow import defaults, serving, TensorFlow, TensorFlowModel, TensorFlowPredictor
 import sagemaker.tensorflow.estimator as tfe
 
 DATA_DIR = os.path.join(os.path.dirname(__file__), "..", "data")
@@ -256,6 +256,7 @@ def test_create_model(sagemaker_session, tf_version):
         container_log_level=container_log_level,
         base_job_name="job",
         source_dir=source_dir,
+        enable_network_isolation=True,
     )
 
     job_name = "doing something"
@@ -271,6 +272,7 @@ def test_create_model(sagemaker_session, tf_version):
     assert model.container_log_level == container_log_level
     assert model.source_dir == source_dir
     assert model.vpc_config is None
+    assert model.enable_network_isolation()
 
 
 def test_create_model_with_optional_params(sagemaker_session):
@@ -1002,11 +1004,23 @@ def test_script_mode_enabled(sagemaker_session):
     assert tf._script_mode_enabled() is False
 
 
-@patch("sagemaker.tensorflow.estimator.TensorFlow._create_tfs_model")
-def test_script_mode_create_model(create_tfs_model, sagemaker_session):
-    tf = _build_tf(sagemaker_session=sagemaker_session, py_version="py3")
-    tf.create_model()
-    create_tfs_model.assert_called_once()
+def test_script_mode_create_model(sagemaker_session):
+    tf = _build_tf(
+        sagemaker_session=sagemaker_session, py_version="py3", enable_network_isolation=True
+    )
+    tf._prepare_for_training()  # set output_path and job name as if training happened
+
+    model = tf.create_model()
+
+    assert isinstance(model, serving.Model)
+
+    assert model.model_data == tf.model_data
+    assert model.role == tf.role
+    assert model.name == tf._current_job_name
+    assert model.container_log_level == tf.container_log_level
+    assert model._framework_version == "1.11"
+    assert model.sagemaker_session == sagemaker_session
+    assert model.enable_network_isolation()
 
 
 @patch("sagemaker.utils.create_tar_file", MagicMock())