feature: network isolation mode for xgboost

src/sagemaker/xgboost/model.py

@@ -145,13 +145,16 @@ def prepare_container_def(self, instance_type=None, accelerator_type=None):
  )
 
  deploy_key_prefix = model_code_key_prefix(self.key_prefix, self.name, deploy_image)
- self._upload_code(deploy_key_prefix)
+ self._upload_code(key_prefix=deploy_key_prefix, repack=self.enable_network_isolation())
  deploy_env = dict(self.env)
  deploy_env.update(self._framework_env_vars())
 
  if self.model_server_workers:
  deploy_env[MODEL_SERVER_WORKERS_PARAM_NAME.upper()] = str(self.model_server_workers)
- return sagemaker.container_def(deploy_image, self.model_data, deploy_env)
+ model_data = (
+ self.repacked_model_data if self.enable_network_isolation() else self.model_data
+ )
+ return sagemaker.container_def(deploy_image, model_data, deploy_env)
 
  def serving_image_uri(self, region_name, instance_type):
  """Create a URI for the serving image.

tests/data/xgboost_abalone/abalone.py

@@ -0,0 +1,50 @@
+import argparse
+import os
+
+from sagemaker_xgboost_container.data_utils import get_dmatrix
+
+import xgboost as xgb
+
+model_filename = "xgboost-model"
+
+if __name__ == "__main__":
+ parser = argparse.ArgumentParser()
+
+ # Sagemaker specific arguments. Defaults are set in the environment variables.
+ parser.add_argument(
+ "--model_dir", type=str, default=os.environ.get("SM_MODEL_DIR", "/opt/ml/model")
+ )
+ parser.add_argument(
+ "--train",
+ type=str,
+ default=os.environ.get("SM_CHANNEL_TRAIN", "/opt/ml/input/data/abalone"),
+ )
+
+ args, _ = parser.parse_known_args()
+
+ dtrain = get_dmatrix(args.train, "libsvm")
+
+ params = {
+ "max_depth": 5,
+ "eta": 0.2,
+ "gamma": 4,
+ "min_child_weight": 6,
+ "subsample": 0.7,
+ "verbosity": 2,
+ "objective": "reg:squarederror",
+ "tree_method": "auto",
+ "predictor": "auto",
+ }
+
+ booster = xgb.train(params=params, dtrain=dtrain, num_boost_round=50)
+ booster.save_model(args.model_dir + "/" + model_filename)
+
+
+def model_fn(model_dir):
+ """Deserialize and return fitted model.
+
+ Note that this should have the same name as the serialized model in the _xgb_train method
+ """
+ booster = xgb.Booster()
+ booster.load_model(os.path.join(model_dir, model_filename))
+ return booster

tests/integ/test_xgboost.py

@@ -14,6 +14,8 @@
 
 import os
 import pytest
+from sagemaker.utils import unique_name_from_base
+from sagemaker.xgboost import XGBoost
 from sagemaker.xgboost.processing import XGBoostProcessor
 from tests.integ import DATA_DIR, TRAINING_DEFAULT_TIMEOUT_MINUTES
 from tests.integ.timeout import timeout
@@ -48,3 +50,35 @@ def test_framework_processing_job_with_deps(
  inputs=[],
  wait=True,
  )
+
+
+def test_training_with_network_isolation(
+ sagemaker_session,
+ xgboost_latest_version,
+ xgboost_latest_py_version,
+ cpu_instance_type,
+):
+ with timeout(minutes=TRAINING_DEFAULT_TIMEOUT_MINUTES):
+ base_job_name = "test-network-isolation-xgboost"
+
+ xgboost = XGBoost(
+ entry_point=os.path.join(DATA_DIR, "xgboost_abalone", "abalone.py"),
+ role=ROLE,
+ instance_type=cpu_instance_type,
+ instance_count=1,
+ framework_version=xgboost_latest_version,
+ py_version=xgboost_latest_py_version,
+ base_job_name=base_job_name,
+ sagemaker_session=sagemaker_session,
+ enable_network_isolation=True,
+ )
+
+ train_input = xgboost.sagemaker_session.upload_data(
+ path=os.path.join(DATA_DIR, "xgboost_abalone", "abalone"),
+ key_prefix="integ-test-data/xgboost_abalone/abalone",
+ )
+ job_name = unique_name_from_base(base_job_name)
+ xgboost.fit(inputs={"train": train_input}, job_name=job_name)
+ assert sagemaker_session.sagemaker_client.describe_training_job(TrainingJobName=job_name)[
+ "EnableNetworkIsolation"
+ ]

tests/unit/test_xgboost.py

@@ -22,6 +22,7 @@
 from packaging.version import Version
 
 
+from sagemaker.fw_utils import UploadedCode
 from sagemaker.xgboost import XGBoost, XGBoostModel, XGBoostPredictor
 
 
@@ -180,6 +181,26 @@ def test_create_model(sagemaker_session, xgboost_framework_version):
  assert model_values["Image"] == default_image_uri
 
 
+@patch("sagemaker.model.FrameworkModel._upload_code")
+def test_create_model_with_network_isolation(upload, sagemaker_session, xgboost_framework_version):
+ source_dir = "s3://mybucket/source"
+ repacked_model_data = "s3://mybucket/prefix/model.tar.gz"
+
+ xgboost_model = XGBoostModel(
+ model_data=source_dir,
+ role=ROLE,
+ sagemaker_session=sagemaker_session,
+ entry_point=SCRIPT_PATH,
+ framework_version=xgboost_framework_version,
+ enable_network_isolation=True,
+ )
+ xgboost_model.uploaded_code = UploadedCode(s3_prefix=repacked_model_data, script_name="script")
+ xgboost_model.repacked_model_data = repacked_model_data
+ model_values = xgboost_model.prepare_container_def(CPU)
+ assert model_values["Environment"]["SAGEMAKER_SUBMIT_DIRECTORY"] == "/opt/ml/model/code"
+ assert model_values["ModelDataUrl"] == repacked_model_data
+
+
 @patch("sagemaker.estimator.name_from_base")
 def test_create_model_from_estimator(name_from_base, sagemaker_session, xgboost_framework_version):
  container_log_level = '"logging.INFO"'

src/sagemaker/xgboost/model.py

@@ -145,13 +145,16 @@ def prepare_container_def(self, instance_type=None, accelerator_type=None):
  )
 
  deploy_key_prefix = model_code_key_prefix(self.key_prefix, self.name, deploy_image)
- self._upload_code(deploy_key_prefix)
+ self._upload_code(key_prefix=deploy_key_prefix, repack=self.enable_network_isolation())
  deploy_env = dict(self.env)
  deploy_env.update(self._framework_env_vars())
 
  if self.model_server_workers:
  deploy_env[MODEL_SERVER_WORKERS_PARAM_NAME.upper()] = str(self.model_server_workers)
- return sagemaker.container_def(deploy_image, self.model_data, deploy_env)
+ model_data = (
+ self.repacked_model_data if self.enable_network_isolation() else self.model_data
+ )
+ return sagemaker.container_def(deploy_image, model_data, deploy_env)
 
  def serving_image_uri(self, region_name, instance_type):
  """Create a URI for the serving image.

tests/data/xgboost_abalone/abalone.py

@@ -0,0 +1,50 @@
+import argparse
+import os
+
+from sagemaker_xgboost_container.data_utils import get_dmatrix
+
+import xgboost as xgb
+
+model_filename = "xgboost-model"
+
+if __name__ == "__main__":
+ parser = argparse.ArgumentParser()
+
+ # Sagemaker specific arguments. Defaults are set in the environment variables.
+ parser.add_argument(
+ "--model_dir", type=str, default=os.environ.get("SM_MODEL_DIR", "/opt/ml/model")
+ )
+ parser.add_argument(
+ "--train",
+ type=str,
+ default=os.environ.get("SM_CHANNEL_TRAIN", "/opt/ml/input/data/abalone"),
+ )
+
+ args, _ = parser.parse_known_args()
+
+ dtrain = get_dmatrix(args.train, "libsvm")
+
+ params = {
+ "max_depth": 5,
+ "eta": 0.2,
+ "gamma": 4,
+ "min_child_weight": 6,
+ "subsample": 0.7,
+ "verbosity": 2,
+ "objective": "reg:squarederror",
+ "tree_method": "auto",
+ "predictor": "auto",
+ }
+
+ booster = xgb.train(params=params, dtrain=dtrain, num_boost_round=50)
+ booster.save_model(args.model_dir + "/" + model_filename)
+
+
+def model_fn(model_dir):
+ """Deserialize and return fitted model.
+
+ Note that this should have the same name as the serialized model in the _xgb_train method
+ """
+ booster = xgb.Booster()
+ booster.load_model(os.path.join(model_dir, model_filename))
+ return booster

tests/integ/test_xgboost.py

@@ -14,6 +14,8 @@
 
 import os
 import pytest
+from sagemaker.utils import unique_name_from_base
+from sagemaker.xgboost import XGBoost
 from sagemaker.xgboost.processing import XGBoostProcessor
 from tests.integ import DATA_DIR, TRAINING_DEFAULT_TIMEOUT_MINUTES
 from tests.integ.timeout import timeout
@@ -48,3 +50,35 @@ def test_framework_processing_job_with_deps(
  inputs=[],
  wait=True,
  )
+
+
+def test_training_with_network_isolation(
+ sagemaker_session,
+ xgboost_latest_version,
+ xgboost_latest_py_version,
+ cpu_instance_type,
+):
+ with timeout(minutes=TRAINING_DEFAULT_TIMEOUT_MINUTES):
+ base_job_name = "test-network-isolation-xgboost"
+
+ xgboost = XGBoost(
+ entry_point=os.path.join(DATA_DIR, "xgboost_abalone", "abalone.py"),
+ role=ROLE,
+ instance_type=cpu_instance_type,
+ instance_count=1,
+ framework_version=xgboost_latest_version,
+ py_version=xgboost_latest_py_version,
+ base_job_name=base_job_name,
+ sagemaker_session=sagemaker_session,
+ enable_network_isolation=True,
+ )
+
+ train_input = xgboost.sagemaker_session.upload_data(
+ path=os.path.join(DATA_DIR, "xgboost_abalone", "abalone"),
+ key_prefix="integ-test-data/xgboost_abalone/abalone",
+ )
+ job_name = unique_name_from_base(base_job_name)
+ xgboost.fit(inputs={"train": train_input}, job_name=job_name)
+ assert sagemaker_session.sagemaker_client.describe_training_job(TrainingJobName=job_name)[
+ "EnableNetworkIsolation"
+ ]

tests/unit/test_xgboost.py

@@ -22,6 +22,7 @@
 from packaging.version import Version
 
 
+from sagemaker.fw_utils import UploadedCode
 from sagemaker.xgboost import XGBoost, XGBoostModel, XGBoostPredictor
 
 
@@ -180,6 +181,26 @@ def test_create_model(sagemaker_session, xgboost_framework_version):
  assert model_values["Image"] == default_image_uri
 
 
+@patch("sagemaker.model.FrameworkModel._upload_code")
+def test_create_model_with_network_isolation(upload, sagemaker_session, xgboost_framework_version):
+ source_dir = "s3://mybucket/source"
+ repacked_model_data = "s3://mybucket/prefix/model.tar.gz"
+
+ xgboost_model = XGBoostModel(
+ model_data=source_dir,
+ role=ROLE,
+ sagemaker_session=sagemaker_session,
+ entry_point=SCRIPT_PATH,
+ framework_version=xgboost_framework_version,
+ enable_network_isolation=True,
+ )
+ xgboost_model.uploaded_code = UploadedCode(s3_prefix=repacked_model_data, script_name="script")
+ xgboost_model.repacked_model_data = repacked_model_data
+ model_values = xgboost_model.prepare_container_def(CPU)
+ assert model_values["Environment"]["SAGEMAKER_SUBMIT_DIRECTORY"] == "/opt/ml/model/code"
+ assert model_values["ModelDataUrl"] == repacked_model_data
+
+
 @patch("sagemaker.estimator.name_from_base")
 def test_create_model_from_estimator(name_from_base, sagemaker_session, xgboost_framework_version):
  container_log_level = '"logging.INFO"'

src/sagemaker/xgboost/model.py

@@ -145,13 +145,16 @@ def prepare_container_def(self, instance_type=None, accelerator_type=None):
  )
 
  deploy_key_prefix = model_code_key_prefix(self.key_prefix, self.name, deploy_image)
- self._upload_code(deploy_key_prefix)
+ self._upload_code(key_prefix=deploy_key_prefix, repack=self.enable_network_isolation())
  deploy_env = dict(self.env)
  deploy_env.update(self._framework_env_vars())
 
  if self.model_server_workers:
  deploy_env[MODEL_SERVER_WORKERS_PARAM_NAME.upper()] = str(self.model_server_workers)
- return sagemaker.container_def(deploy_image, self.model_data, deploy_env)
+ model_data = (
+ self.repacked_model_data if self.enable_network_isolation() else self.model_data
+ )
+ return sagemaker.container_def(deploy_image, model_data, deploy_env)
 
  def serving_image_uri(self, region_name, instance_type):
  """Create a URI for the serving image.

tests/data/xgboost_abalone/abalone.py

@@ -0,0 +1,50 @@
+import argparse
+import os
+
+from sagemaker_xgboost_container.data_utils import get_dmatrix
+
+import xgboost as xgb
+
+model_filename = "xgboost-model"
+
+if __name__ == "__main__":
+ parser = argparse.ArgumentParser()
+
+ # Sagemaker specific arguments. Defaults are set in the environment variables.
+ parser.add_argument(
+ "--model_dir", type=str, default=os.environ.get("SM_MODEL_DIR", "/opt/ml/model")
+ )
+ parser.add_argument(
+ "--train",
+ type=str,
+ default=os.environ.get("SM_CHANNEL_TRAIN", "/opt/ml/input/data/abalone"),
+ )
+
+ args, _ = parser.parse_known_args()
+
+ dtrain = get_dmatrix(args.train, "libsvm")
+
+ params = {
+ "max_depth": 5,
+ "eta": 0.2,
+ "gamma": 4,
+ "min_child_weight": 6,
+ "subsample": 0.7,
+ "verbosity": 2,
+ "objective": "reg:squarederror",
+ "tree_method": "auto",
+ "predictor": "auto",
+ }
+
+ booster = xgb.train(params=params, dtrain=dtrain, num_boost_round=50)
+ booster.save_model(args.model_dir + "/" + model_filename)
+
+
+def model_fn(model_dir):
+ """Deserialize and return fitted model.
+
+ Note that this should have the same name as the serialized model in the _xgb_train method
+ """
+ booster = xgb.Booster()
+ booster.load_model(os.path.join(model_dir, model_filename))
+ return booster

tests/integ/test_xgboost.py

@@ -14,6 +14,8 @@
 
 import os
 import pytest
+from sagemaker.utils import unique_name_from_base
+from sagemaker.xgboost import XGBoost
 from sagemaker.xgboost.processing import XGBoostProcessor
 from tests.integ import DATA_DIR, TRAINING_DEFAULT_TIMEOUT_MINUTES
 from tests.integ.timeout import timeout
@@ -48,3 +50,35 @@ def test_framework_processing_job_with_deps(
  inputs=[],
  wait=True,
  )
+
+
+def test_training_with_network_isolation(
+ sagemaker_session,
+ xgboost_latest_version,
+ xgboost_latest_py_version,
+ cpu_instance_type,
+):
+ with timeout(minutes=TRAINING_DEFAULT_TIMEOUT_MINUTES):
+ base_job_name = "test-network-isolation-xgboost"
+
+ xgboost = XGBoost(
+ entry_point=os.path.join(DATA_DIR, "xgboost_abalone", "abalone.py"),
+ role=ROLE,
+ instance_type=cpu_instance_type,
+ instance_count=1,
+ framework_version=xgboost_latest_version,
+ py_version=xgboost_latest_py_version,
+ base_job_name=base_job_name,
+ sagemaker_session=sagemaker_session,
+ enable_network_isolation=True,
+ )
+
+ train_input = xgboost.sagemaker_session.upload_data(
+ path=os.path.join(DATA_DIR, "xgboost_abalone", "abalone"),
+ key_prefix="integ-test-data/xgboost_abalone/abalone",
+ )
+ job_name = unique_name_from_base(base_job_name)
+ xgboost.fit(inputs={"train": train_input}, job_name=job_name)
+ assert sagemaker_session.sagemaker_client.describe_training_job(TrainingJobName=job_name)[
+ "EnableNetworkIsolation"
+ ]

tests/unit/test_xgboost.py

@@ -22,6 +22,7 @@
 from packaging.version import Version
 
 
+from sagemaker.fw_utils import UploadedCode
 from sagemaker.xgboost import XGBoost, XGBoostModel, XGBoostPredictor
 
 
@@ -180,6 +181,26 @@ def test_create_model(sagemaker_session, xgboost_framework_version):
  assert model_values["Image"] == default_image_uri
 
 
+@patch("sagemaker.model.FrameworkModel._upload_code")
+def test_create_model_with_network_isolation(upload, sagemaker_session, xgboost_framework_version):
+ source_dir = "s3://mybucket/source"
+ repacked_model_data = "s3://mybucket/prefix/model.tar.gz"
+
+ xgboost_model = XGBoostModel(
+ model_data=source_dir,
+ role=ROLE,
+ sagemaker_session=sagemaker_session,
+ entry_point=SCRIPT_PATH,
+ framework_version=xgboost_framework_version,
+ enable_network_isolation=True,
+ )
+ xgboost_model.uploaded_code = UploadedCode(s3_prefix=repacked_model_data, script_name="script")
+ xgboost_model.repacked_model_data = repacked_model_data
+ model_values = xgboost_model.prepare_container_def(CPU)
+ assert model_values["Environment"]["SAGEMAKER_SUBMIT_DIRECTORY"] == "/opt/ml/model/code"
+ assert model_values["ModelDataUrl"] == repacked_model_data
+
+
 @patch("sagemaker.estimator.name_from_base")
 def test_create_model_from_estimator(name_from_base, sagemaker_session, xgboost_framework_version):
  container_log_level = '"logging.INFO"'

src/sagemaker/xgboost/model.py

@@ -145,13 +145,16 @@ def prepare_container_def(self, instance_type=None, accelerator_type=None):
  )
 
  deploy_key_prefix = model_code_key_prefix(self.key_prefix, self.name, deploy_image)
- self._upload_code(deploy_key_prefix)
+ self._upload_code(key_prefix=deploy_key_prefix, repack=self.enable_network_isolation())
  deploy_env = dict(self.env)
  deploy_env.update(self._framework_env_vars())
 
  if self.model_server_workers:
  deploy_env[MODEL_SERVER_WORKERS_PARAM_NAME.upper()] = str(self.model_server_workers)
- return sagemaker.container_def(deploy_image, self.model_data, deploy_env)
+ model_data = (
+ self.repacked_model_data if self.enable_network_isolation() else self.model_data
+ )
+ return sagemaker.container_def(deploy_image, model_data, deploy_env)
 
  def serving_image_uri(self, region_name, instance_type):
  """Create a URI for the serving image.

tests/data/xgboost_abalone/abalone.py

@@ -0,0 +1,50 @@
+import argparse
+import os
+
+from sagemaker_xgboost_container.data_utils import get_dmatrix
+
+import xgboost as xgb
+
+model_filename = "xgboost-model"
+
+if __name__ == "__main__":
+ parser = argparse.ArgumentParser()
+
+ # Sagemaker specific arguments. Defaults are set in the environment variables.
+ parser.add_argument(
+ "--model_dir", type=str, default=os.environ.get("SM_MODEL_DIR", "/opt/ml/model")
+ )
+ parser.add_argument(
+ "--train",
+ type=str,
+ default=os.environ.get("SM_CHANNEL_TRAIN", "/opt/ml/input/data/abalone"),
+ )
+
+ args, _ = parser.parse_known_args()
+
+ dtrain = get_dmatrix(args.train, "libsvm")
+
+ params = {
+ "max_depth": 5,
+ "eta": 0.2,
+ "gamma": 4,
+ "min_child_weight": 6,
+ "subsample": 0.7,
+ "verbosity": 2,
+ "objective": "reg:squarederror",
+ "tree_method": "auto",
+ "predictor": "auto",
+ }
+
+ booster = xgb.train(params=params, dtrain=dtrain, num_boost_round=50)
+ booster.save_model(args.model_dir + "/" + model_filename)
+
+
+def model_fn(model_dir):
+ """Deserialize and return fitted model.
+
+ Note that this should have the same name as the serialized model in the _xgb_train method
+ """
+ booster = xgb.Booster()
+ booster.load_model(os.path.join(model_dir, model_filename))
+ return booster

tests/integ/test_xgboost.py

@@ -14,6 +14,8 @@
 
 import os
 import pytest
+from sagemaker.utils import unique_name_from_base
+from sagemaker.xgboost import XGBoost
 from sagemaker.xgboost.processing import XGBoostProcessor
 from tests.integ import DATA_DIR, TRAINING_DEFAULT_TIMEOUT_MINUTES
 from tests.integ.timeout import timeout
@@ -48,3 +50,35 @@ def test_framework_processing_job_with_deps(
  inputs=[],
  wait=True,
  )
+
+
+def test_training_with_network_isolation(
+ sagemaker_session,
+ xgboost_latest_version,
+ xgboost_latest_py_version,
+ cpu_instance_type,
+):
+ with timeout(minutes=TRAINING_DEFAULT_TIMEOUT_MINUTES):
+ base_job_name = "test-network-isolation-xgboost"
+
+ xgboost = XGBoost(
+ entry_point=os.path.join(DATA_DIR, "xgboost_abalone", "abalone.py"),
+ role=ROLE,
+ instance_type=cpu_instance_type,
+ instance_count=1,
+ framework_version=xgboost_latest_version,
+ py_version=xgboost_latest_py_version,
+ base_job_name=base_job_name,
+ sagemaker_session=sagemaker_session,
+ enable_network_isolation=True,
+ )
+
+ train_input = xgboost.sagemaker_session.upload_data(
+ path=os.path.join(DATA_DIR, "xgboost_abalone", "abalone"),
+ key_prefix="integ-test-data/xgboost_abalone/abalone",
+ )
+ job_name = unique_name_from_base(base_job_name)
+ xgboost.fit(inputs={"train": train_input}, job_name=job_name)
+ assert sagemaker_session.sagemaker_client.describe_training_job(TrainingJobName=job_name)[
+ "EnableNetworkIsolation"
+ ]

tests/unit/test_xgboost.py

@@ -22,6 +22,7 @@
 from packaging.version import Version
 
 
+from sagemaker.fw_utils import UploadedCode
 from sagemaker.xgboost import XGBoost, XGBoostModel, XGBoostPredictor
 
 
@@ -180,6 +181,26 @@ def test_create_model(sagemaker_session, xgboost_framework_version):
  assert model_values["Image"] == default_image_uri
 
 
+@patch("sagemaker.model.FrameworkModel._upload_code")
+def test_create_model_with_network_isolation(upload, sagemaker_session, xgboost_framework_version):
+ source_dir = "s3://mybucket/source"
+ repacked_model_data = "s3://mybucket/prefix/model.tar.gz"
+
+ xgboost_model = XGBoostModel(
+ model_data=source_dir,
+ role=ROLE,
+ sagemaker_session=sagemaker_session,
+ entry_point=SCRIPT_PATH,
+ framework_version=xgboost_framework_version,
+ enable_network_isolation=True,
+ )
+ xgboost_model.uploaded_code = UploadedCode(s3_prefix=repacked_model_data, script_name="script")
+ xgboost_model.repacked_model_data = repacked_model_data
+ model_values = xgboost_model.prepare_container_def(CPU)
+ assert model_values["Environment"]["SAGEMAKER_SUBMIT_DIRECTORY"] == "/opt/ml/model/code"
+ assert model_values["ModelDataUrl"] == repacked_model_data
+
+
 @patch("sagemaker.estimator.name_from_base")
 def test_create_model_from_estimator(name_from_base, sagemaker_session, xgboost_framework_version):
  container_log_level = '"logging.INFO"'

src/sagemaker/xgboost/model.py

@@ -145,13 +145,16 @@ def prepare_container_def(self, instance_type=None, accelerator_type=None):
  )
 
  deploy_key_prefix = model_code_key_prefix(self.key_prefix, self.name, deploy_image)
- self._upload_code(deploy_key_prefix)
+ self._upload_code(key_prefix=deploy_key_prefix, repack=self.enable_network_isolation())
  deploy_env = dict(self.env)
  deploy_env.update(self._framework_env_vars())
 
  if self.model_server_workers:
  deploy_env[MODEL_SERVER_WORKERS_PARAM_NAME.upper()] = str(self.model_server_workers)
- return sagemaker.container_def(deploy_image, self.model_data, deploy_env)
+ model_data = (
+ self.repacked_model_data if self.enable_network_isolation() else self.model_data
+ )
+ return sagemaker.container_def(deploy_image, model_data, deploy_env)
 
  def serving_image_uri(self, region_name, instance_type):
  """Create a URI for the serving image.

tests/data/xgboost_abalone/abalone.py

@@ -0,0 +1,50 @@
+import argparse
+import os
+
+from sagemaker_xgboost_container.data_utils import get_dmatrix
+
+import xgboost as xgb
+
+model_filename = "xgboost-model"
+
+if __name__ == "__main__":
+ parser = argparse.ArgumentParser()
+
+ # Sagemaker specific arguments. Defaults are set in the environment variables.
+ parser.add_argument(
+ "--model_dir", type=str, default=os.environ.get("SM_MODEL_DIR", "/opt/ml/model")
+ )
+ parser.add_argument(
+ "--train",
+ type=str,
+ default=os.environ.get("SM_CHANNEL_TRAIN", "/opt/ml/input/data/abalone"),
+ )
+
+ args, _ = parser.parse_known_args()
+
+ dtrain = get_dmatrix(args.train, "libsvm")
+
+ params = {
+ "max_depth": 5,
+ "eta": 0.2,
+ "gamma": 4,
+ "min_child_weight": 6,
+ "subsample": 0.7,
+ "verbosity": 2,
+ "objective": "reg:squarederror",
+ "tree_method": "auto",
+ "predictor": "auto",
+ }
+
+ booster = xgb.train(params=params, dtrain=dtrain, num_boost_round=50)
+ booster.save_model(args.model_dir + "/" + model_filename)
+
+
+def model_fn(model_dir):
+ """Deserialize and return fitted model.
+
+ Note that this should have the same name as the serialized model in the _xgb_train method
+ """
+ booster = xgb.Booster()
+ booster.load_model(os.path.join(model_dir, model_filename))
+ return booster

tests/integ/test_xgboost.py

@@ -14,6 +14,8 @@
 
 import os
 import pytest
+from sagemaker.utils import unique_name_from_base
+from sagemaker.xgboost import XGBoost
 from sagemaker.xgboost.processing import XGBoostProcessor
 from tests.integ import DATA_DIR, TRAINING_DEFAULT_TIMEOUT_MINUTES
 from tests.integ.timeout import timeout
@@ -48,3 +50,35 @@ def test_framework_processing_job_with_deps(
  inputs=[],
  wait=True,
  )
+
+
+def test_training_with_network_isolation(
+ sagemaker_session,
+ xgboost_latest_version,
+ xgboost_latest_py_version,
+ cpu_instance_type,
+):
+ with timeout(minutes=TRAINING_DEFAULT_TIMEOUT_MINUTES):
+ base_job_name = "test-network-isolation-xgboost"
+
+ xgboost = XGBoost(
+ entry_point=os.path.join(DATA_DIR, "xgboost_abalone", "abalone.py"),
+ role=ROLE,
+ instance_type=cpu_instance_type,
+ instance_count=1,
+ framework_version=xgboost_latest_version,
+ py_version=xgboost_latest_py_version,
+ base_job_name=base_job_name,
+ sagemaker_session=sagemaker_session,
+ enable_network_isolation=True,
+ )
+
+ train_input = xgboost.sagemaker_session.upload_data(
+ path=os.path.join(DATA_DIR, "xgboost_abalone", "abalone"),
+ key_prefix="integ-test-data/xgboost_abalone/abalone",
+ )
+ job_name = unique_name_from_base(base_job_name)
+ xgboost.fit(inputs={"train": train_input}, job_name=job_name)
+ assert sagemaker_session.sagemaker_client.describe_training_job(TrainingJobName=job_name)[
+ "EnableNetworkIsolation"
+ ]

tests/unit/test_xgboost.py

@@ -22,6 +22,7 @@
 from packaging.version import Version
 
 
+from sagemaker.fw_utils import UploadedCode
 from sagemaker.xgboost import XGBoost, XGBoostModel, XGBoostPredictor
 
 
@@ -180,6 +181,26 @@ def test_create_model(sagemaker_session, xgboost_framework_version):
  assert model_values["Image"] == default_image_uri
 
 
+@patch("sagemaker.model.FrameworkModel._upload_code")
+def test_create_model_with_network_isolation(upload, sagemaker_session, xgboost_framework_version):
+ source_dir = "s3://mybucket/source"
+ repacked_model_data = "s3://mybucket/prefix/model.tar.gz"
+
+ xgboost_model = XGBoostModel(
+ model_data=source_dir,
+ role=ROLE,
+ sagemaker_session=sagemaker_session,
+ entry_point=SCRIPT_PATH,
+ framework_version=xgboost_framework_version,
+ enable_network_isolation=True,
+ )
+ xgboost_model.uploaded_code = UploadedCode(s3_prefix=repacked_model_data, script_name="script")
+ xgboost_model.repacked_model_data = repacked_model_data
+ model_values = xgboost_model.prepare_container_def(CPU)
+ assert model_values["Environment"]["SAGEMAKER_SUBMIT_DIRECTORY"] == "/opt/ml/model/code"
+ assert model_values["ModelDataUrl"] == repacked_model_data
+
+
 @patch("sagemaker.estimator.name_from_base")
 def test_create_model_from_estimator(name_from_base, sagemaker_session, xgboost_framework_version):
  container_log_level = '"logging.INFO"'

src/sagemaker/xgboost/model.py

@@ -145,13 +145,16 @@ def prepare_container_def(self, instance_type=None, accelerator_type=None):
  )
 
  deploy_key_prefix = model_code_key_prefix(self.key_prefix, self.name, deploy_image)
- self._upload_code(deploy_key_prefix)
+ self._upload_code(key_prefix=deploy_key_prefix, repack=self.enable_network_isolation())
  deploy_env = dict(self.env)
  deploy_env.update(self._framework_env_vars())
 
  if self.model_server_workers:
  deploy_env[MODEL_SERVER_WORKERS_PARAM_NAME.upper()] = str(self.model_server_workers)
- return sagemaker.container_def(deploy_image, self.model_data, deploy_env)
+ model_data = (
+ self.repacked_model_data if self.enable_network_isolation() else self.model_data
+ )
+ return sagemaker.container_def(deploy_image, model_data, deploy_env)
 
  def serving_image_uri(self, region_name, instance_type):
  """Create a URI for the serving image.

tests/data/xgboost_abalone/abalone.py

@@ -0,0 +1,50 @@
+import argparse
+import os
+
+from sagemaker_xgboost_container.data_utils import get_dmatrix
+
+import xgboost as xgb
+
+model_filename = "xgboost-model"
+
+if __name__ == "__main__":
+ parser = argparse.ArgumentParser()
+
+ # Sagemaker specific arguments. Defaults are set in the environment variables.
+ parser.add_argument(
+ "--model_dir", type=str, default=os.environ.get("SM_MODEL_DIR", "/opt/ml/model")
+ )
+ parser.add_argument(
+ "--train",
+ type=str,
+ default=os.environ.get("SM_CHANNEL_TRAIN", "/opt/ml/input/data/abalone"),
+ )
+
+ args, _ = parser.parse_known_args()
+
+ dtrain = get_dmatrix(args.train, "libsvm")
+
+ params = {
+ "max_depth": 5,
+ "eta": 0.2,
+ "gamma": 4,
+ "min_child_weight": 6,
+ "subsample": 0.7,
+ "verbosity": 2,
+ "objective": "reg:squarederror",
+ "tree_method": "auto",
+ "predictor": "auto",
+ }
+
+ booster = xgb.train(params=params, dtrain=dtrain, num_boost_round=50)
+ booster.save_model(args.model_dir + "/" + model_filename)
+
+
+def model_fn(model_dir):
+ """Deserialize and return fitted model.
+
+ Note that this should have the same name as the serialized model in the _xgb_train method
+ """
+ booster = xgb.Booster()
+ booster.load_model(os.path.join(model_dir, model_filename))
+ return booster

tests/integ/test_xgboost.py

@@ -14,6 +14,8 @@
 
 import os
 import pytest
+from sagemaker.utils import unique_name_from_base
+from sagemaker.xgboost import XGBoost
 from sagemaker.xgboost.processing import XGBoostProcessor
 from tests.integ import DATA_DIR, TRAINING_DEFAULT_TIMEOUT_MINUTES
 from tests.integ.timeout import timeout
@@ -48,3 +50,35 @@ def test_framework_processing_job_with_deps(
  inputs=[],
  wait=True,
  )
+
+
+def test_training_with_network_isolation(
+ sagemaker_session,
+ xgboost_latest_version,
+ xgboost_latest_py_version,
+ cpu_instance_type,
+):
+ with timeout(minutes=TRAINING_DEFAULT_TIMEOUT_MINUTES):
+ base_job_name = "test-network-isolation-xgboost"
+
+ xgboost = XGBoost(
+ entry_point=os.path.join(DATA_DIR, "xgboost_abalone", "abalone.py"),
+ role=ROLE,
+ instance_type=cpu_instance_type,
+ instance_count=1,
+ framework_version=xgboost_latest_version,
+ py_version=xgboost_latest_py_version,
+ base_job_name=base_job_name,
+ sagemaker_session=sagemaker_session,
+ enable_network_isolation=True,
+ )
+
+ train_input = xgboost.sagemaker_session.upload_data(
+ path=os.path.join(DATA_DIR, "xgboost_abalone", "abalone"),
+ key_prefix="integ-test-data/xgboost_abalone/abalone",
+ )
+ job_name = unique_name_from_base(base_job_name)
+ xgboost.fit(inputs={"train": train_input}, job_name=job_name)
+ assert sagemaker_session.sagemaker_client.describe_training_job(TrainingJobName=job_name)[
+ "EnableNetworkIsolation"
+ ]

tests/unit/test_xgboost.py

@@ -22,6 +22,7 @@
 from packaging.version import Version
 
 
+from sagemaker.fw_utils import UploadedCode
 from sagemaker.xgboost import XGBoost, XGBoostModel, XGBoostPredictor
 
 
@@ -180,6 +181,26 @@ def test_create_model(sagemaker_session, xgboost_framework_version):
  assert model_values["Image"] == default_image_uri
 
 
+@patch("sagemaker.model.FrameworkModel._upload_code")
+def test_create_model_with_network_isolation(upload, sagemaker_session, xgboost_framework_version):
+ source_dir = "s3://mybucket/source"
+ repacked_model_data = "s3://mybucket/prefix/model.tar.gz"
+
+ xgboost_model = XGBoostModel(
+ model_data=source_dir,
+ role=ROLE,
+ sagemaker_session=sagemaker_session,
+ entry_point=SCRIPT_PATH,
+ framework_version=xgboost_framework_version,
+ enable_network_isolation=True,
+ )
+ xgboost_model.uploaded_code = UploadedCode(s3_prefix=repacked_model_data, script_name="script")
+ xgboost_model.repacked_model_data = repacked_model_data
+ model_values = xgboost_model.prepare_container_def(CPU)
+ assert model_values["Environment"]["SAGEMAKER_SUBMIT_DIRECTORY"] == "/opt/ml/model/code"
+ assert model_values["ModelDataUrl"] == repacked_model_data
+
+
 @patch("sagemaker.estimator.name_from_base")
 def test_create_model_from_estimator(name_from_base, sagemaker_session, xgboost_framework_version):
  container_log_level = '"logging.INFO"'

src/sagemaker/xgboost/model.py

@@ -145,13 +145,16 @@ def prepare_container_def(self, instance_type=None, accelerator_type=None):
  )
 
  deploy_key_prefix = model_code_key_prefix(self.key_prefix, self.name, deploy_image)
- self._upload_code(deploy_key_prefix)
+ self._upload_code(key_prefix=deploy_key_prefix, repack=self.enable_network_isolation())
  deploy_env = dict(self.env)
  deploy_env.update(self._framework_env_vars())
 
  if self.model_server_workers:
  deploy_env[MODEL_SERVER_WORKERS_PARAM_NAME.upper()] = str(self.model_server_workers)
- return sagemaker.container_def(deploy_image, self.model_data, deploy_env)
+ model_data = (
+ self.repacked_model_data if self.enable_network_isolation() else self.model_data
+ )
+ return sagemaker.container_def(deploy_image, model_data, deploy_env)
 
  def serving_image_uri(self, region_name, instance_type):
  """Create a URI for the serving image.

tests/data/xgboost_abalone/abalone.py

@@ -0,0 +1,50 @@
+import argparse
+import os
+
+from sagemaker_xgboost_container.data_utils import get_dmatrix
+
+import xgboost as xgb
+
+model_filename = "xgboost-model"
+
+if __name__ == "__main__":
+ parser = argparse.ArgumentParser()
+
+ # Sagemaker specific arguments. Defaults are set in the environment variables.
+ parser.add_argument(
+ "--model_dir", type=str, default=os.environ.get("SM_MODEL_DIR", "/opt/ml/model")
+ )
+ parser.add_argument(
+ "--train",
+ type=str,
+ default=os.environ.get("SM_CHANNEL_TRAIN", "/opt/ml/input/data/abalone"),
+ )
+
+ args, _ = parser.parse_known_args()
+
+ dtrain = get_dmatrix(args.train, "libsvm")
+
+ params = {
+ "max_depth": 5,
+ "eta": 0.2,
+ "gamma": 4,
+ "min_child_weight": 6,
+ "subsample": 0.7,
+ "verbosity": 2,
+ "objective": "reg:squarederror",
+ "tree_method": "auto",
+ "predictor": "auto",
+ }
+
+ booster = xgb.train(params=params, dtrain=dtrain, num_boost_round=50)
+ booster.save_model(args.model_dir + "/" + model_filename)
+
+
+def model_fn(model_dir):
+ """Deserialize and return fitted model.
+
+ Note that this should have the same name as the serialized model in the _xgb_train method
+ """
+ booster = xgb.Booster()
+ booster.load_model(os.path.join(model_dir, model_filename))
+ return booster

tests/integ/test_xgboost.py

@@ -14,6 +14,8 @@
 
 import os
 import pytest
+from sagemaker.utils import unique_name_from_base
+from sagemaker.xgboost import XGBoost
 from sagemaker.xgboost.processing import XGBoostProcessor
 from tests.integ import DATA_DIR, TRAINING_DEFAULT_TIMEOUT_MINUTES
 from tests.integ.timeout import timeout
@@ -48,3 +50,35 @@ def test_framework_processing_job_with_deps(
  inputs=[],
  wait=True,
  )
+
+
+def test_training_with_network_isolation(
+ sagemaker_session,
+ xgboost_latest_version,
+ xgboost_latest_py_version,
+ cpu_instance_type,
+):
+ with timeout(minutes=TRAINING_DEFAULT_TIMEOUT_MINUTES):
+ base_job_name = "test-network-isolation-xgboost"
+
+ xgboost = XGBoost(
+ entry_point=os.path.join(DATA_DIR, "xgboost_abalone", "abalone.py"),
+ role=ROLE,
+ instance_type=cpu_instance_type,
+ instance_count=1,
+ framework_version=xgboost_latest_version,
+ py_version=xgboost_latest_py_version,
+ base_job_name=base_job_name,
+ sagemaker_session=sagemaker_session,
+ enable_network_isolation=True,
+ )
+
+ train_input = xgboost.sagemaker_session.upload_data(
+ path=os.path.join(DATA_DIR, "xgboost_abalone", "abalone"),
+ key_prefix="integ-test-data/xgboost_abalone/abalone",
+ )
+ job_name = unique_name_from_base(base_job_name)
+ xgboost.fit(inputs={"train": train_input}, job_name=job_name)
+ assert sagemaker_session.sagemaker_client.describe_training_job(TrainingJobName=job_name)[
+ "EnableNetworkIsolation"
+ ]

tests/unit/test_xgboost.py

@@ -22,6 +22,7 @@
 from packaging.version import Version
 
 
+from sagemaker.fw_utils import UploadedCode
 from sagemaker.xgboost import XGBoost, XGBoostModel, XGBoostPredictor
 
 
@@ -180,6 +181,26 @@ def test_create_model(sagemaker_session, xgboost_framework_version):
  assert model_values["Image"] == default_image_uri
 
 
+@patch("sagemaker.model.FrameworkModel._upload_code")
+def test_create_model_with_network_isolation(upload, sagemaker_session, xgboost_framework_version):
+ source_dir = "s3://mybucket/source"
+ repacked_model_data = "s3://mybucket/prefix/model.tar.gz"
+
+ xgboost_model = XGBoostModel(
+ model_data=source_dir,
+ role=ROLE,
+ sagemaker_session=sagemaker_session,
+ entry_point=SCRIPT_PATH,
+ framework_version=xgboost_framework_version,
+ enable_network_isolation=True,
+ )
+ xgboost_model.uploaded_code = UploadedCode(s3_prefix=repacked_model_data, script_name="script")
+ xgboost_model.repacked_model_data = repacked_model_data
+ model_values = xgboost_model.prepare_container_def(CPU)
+ assert model_values["Environment"]["SAGEMAKER_SUBMIT_DIRECTORY"] == "/opt/ml/model/code"
+ assert model_values["ModelDataUrl"] == repacked_model_data
+
+
 @patch("sagemaker.estimator.name_from_base")
 def test_create_model_from_estimator(name_from_base, sagemaker_session, xgboost_framework_version):
  container_log_level = '"logging.INFO"'

src/sagemaker/xgboost/model.py

@@ -145,13 +145,16 @@ def prepare_container_def(self, instance_type=None, accelerator_type=None):
  )
 
  deploy_key_prefix = model_code_key_prefix(self.key_prefix, self.name, deploy_image)
- self._upload_code(deploy_key_prefix)
+ self._upload_code(key_prefix=deploy_key_prefix, repack=self.enable_network_isolation())
  deploy_env = dict(self.env)
  deploy_env.update(self._framework_env_vars())
 
  if self.model_server_workers:
  deploy_env[MODEL_SERVER_WORKERS_PARAM_NAME.upper()] = str(self.model_server_workers)
- return sagemaker.container_def(deploy_image, self.model_data, deploy_env)
+ model_data = (
+ self.repacked_model_data if self.enable_network_isolation() else self.model_data
+ )
+ return sagemaker.container_def(deploy_image, model_data, deploy_env)
 
  def serving_image_uri(self, region_name, instance_type):
  """Create a URI for the serving image.

tests/data/xgboost_abalone/abalone.py

@@ -0,0 +1,50 @@
+import argparse
+import os
+
+from sagemaker_xgboost_container.data_utils import get_dmatrix
+
+import xgboost as xgb
+
+model_filename = "xgboost-model"
+
+if __name__ == "__main__":
+ parser = argparse.ArgumentParser()
+
+ # Sagemaker specific arguments. Defaults are set in the environment variables.
+ parser.add_argument(
+ "--model_dir", type=str, default=os.environ.get("SM_MODEL_DIR", "/opt/ml/model")
+ )
+ parser.add_argument(
+ "--train",
+ type=str,
+ default=os.environ.get("SM_CHANNEL_TRAIN", "/opt/ml/input/data/abalone"),
+ )
+
+ args, _ = parser.parse_known_args()
+
+ dtrain = get_dmatrix(args.train, "libsvm")
+
+ params = {
+ "max_depth": 5,
+ "eta": 0.2,
+ "gamma": 4,
+ "min_child_weight": 6,
+ "subsample": 0.7,
+ "verbosity": 2,
+ "objective": "reg:squarederror",
+ "tree_method": "auto",
+ "predictor": "auto",
+ }
+
+ booster = xgb.train(params=params, dtrain=dtrain, num_boost_round=50)
+ booster.save_model(args.model_dir + "/" + model_filename)
+
+
+def model_fn(model_dir):
+ """Deserialize and return fitted model.
+
+ Note that this should have the same name as the serialized model in the _xgb_train method
+ """
+ booster = xgb.Booster()
+ booster.load_model(os.path.join(model_dir, model_filename))
+ return booster

tests/integ/test_xgboost.py

@@ -14,6 +14,8 @@
 
 import os
 import pytest
+from sagemaker.utils import unique_name_from_base
+from sagemaker.xgboost import XGBoost
 from sagemaker.xgboost.processing import XGBoostProcessor
 from tests.integ import DATA_DIR, TRAINING_DEFAULT_TIMEOUT_MINUTES
 from tests.integ.timeout import timeout
@@ -48,3 +50,35 @@ def test_framework_processing_job_with_deps(
  inputs=[],
  wait=True,
  )
+
+
+def test_training_with_network_isolation(
+ sagemaker_session,
+ xgboost_latest_version,
+ xgboost_latest_py_version,
+ cpu_instance_type,
+):
+ with timeout(minutes=TRAINING_DEFAULT_TIMEOUT_MINUTES):
+ base_job_name = "test-network-isolation-xgboost"
+
+ xgboost = XGBoost(
+ entry_point=os.path.join(DATA_DIR, "xgboost_abalone", "abalone.py"),
+ role=ROLE,
+ instance_type=cpu_instance_type,
+ instance_count=1,
+ framework_version=xgboost_latest_version,
+ py_version=xgboost_latest_py_version,
+ base_job_name=base_job_name,
+ sagemaker_session=sagemaker_session,
+ enable_network_isolation=True,
+ )
+
+ train_input = xgboost.sagemaker_session.upload_data(
+ path=os.path.join(DATA_DIR, "xgboost_abalone", "abalone"),
+ key_prefix="integ-test-data/xgboost_abalone/abalone",
+ )
+ job_name = unique_name_from_base(base_job_name)
+ xgboost.fit(inputs={"train": train_input}, job_name=job_name)
+ assert sagemaker_session.sagemaker_client.describe_training_job(TrainingJobName=job_name)[
+ "EnableNetworkIsolation"
+ ]

tests/unit/test_xgboost.py

@@ -22,6 +22,7 @@
 from packaging.version import Version
 
 
+from sagemaker.fw_utils import UploadedCode
 from sagemaker.xgboost import XGBoost, XGBoostModel, XGBoostPredictor
 
 
@@ -180,6 +181,26 @@ def test_create_model(sagemaker_session, xgboost_framework_version):
  assert model_values["Image"] == default_image_uri
 
 
+@patch("sagemaker.model.FrameworkModel._upload_code")
+def test_create_model_with_network_isolation(upload, sagemaker_session, xgboost_framework_version):
+ source_dir = "s3://mybucket/source"
+ repacked_model_data = "s3://mybucket/prefix/model.tar.gz"
+
+ xgboost_model = XGBoostModel(
+ model_data=source_dir,
+ role=ROLE,
+ sagemaker_session=sagemaker_session,
+ entry_point=SCRIPT_PATH,
+ framework_version=xgboost_framework_version,
+ enable_network_isolation=True,
+ )
+ xgboost_model.uploaded_code = UploadedCode(s3_prefix=repacked_model_data, script_name="script")
+ xgboost_model.repacked_model_data = repacked_model_data
+ model_values = xgboost_model.prepare_container_def(CPU)
+ assert model_values["Environment"]["SAGEMAKER_SUBMIT_DIRECTORY"] == "/opt/ml/model/code"
+ assert model_values["ModelDataUrl"] == repacked_model_data
+
+
 @patch("sagemaker.estimator.name_from_base")
 def test_create_model_from_estimator(name_from_base, sagemaker_session, xgboost_framework_version):
  container_log_level = '"logging.INFO"'