From a58ce71ae64c80f340a5e93364ec1dc1da99ce35 Mon Sep 17 00:00:00 2001 From: Brent Millare Date: Fri, 10 Sep 2021 20:40:09 +0000 Subject: [PATCH 1/7] feature: network isolation mode for xgboost --- src/sagemaker/xgboost/model.py | 7 +++++-- tests/unit/test_xgboost.py | 21 +++++++++++++++++++++ 2 files changed, 26 insertions(+), 2 deletions(-) diff --git a/src/sagemaker/xgboost/model.py b/src/sagemaker/xgboost/model.py index ae50d76329..49acc11074 100644 --- a/src/sagemaker/xgboost/model.py +++ b/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. diff --git a/tests/unit/test_xgboost.py b/tests/unit/test_xgboost.py index 28937ccba1..82f27c19ae 100644 --- a/tests/unit/test_xgboost.py +++ b/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"' From b038943a89abefd3503d97c297cd21e0b495dce1 Mon Sep 17 00:00:00 2001 From: Brent Millare Date: Fri, 10 Sep 2021 20:35:13 +0000 Subject: [PATCH 2/7] feature: network isolation mode for xgboost --- tox.ini | 1 + 1 file changed, 1 insertion(+) diff --git a/tox.ini b/tox.ini index d1a503be88..57d0d88b00 100644 --- a/tox.ini +++ b/tox.ini @@ -58,6 +58,7 @@ markers = timeout: mark a test as a timeout. [testenv] +basepython = /usr/bin/python3 passenv = AWS_ACCESS_KEY_ID AWS_SECRET_ACCESS_KEY From 77a5a5e1272b82ba933102907556170272fca2da Mon Sep 17 00:00:00 2001 From: Brent Millare Date: Fri, 8 Oct 2021 20:25:58 +0000 Subject: [PATCH 3/7] Add integ tests for xgboost net iso --- tests/data/xgboost_abalone/abalone.py | 46 +++++++++++++++++++++++++++ tests/integ/test_xgboost.py | 35 ++++++++++++++++++++ 2 files changed, 81 insertions(+) create mode 100644 tests/data/xgboost_abalone/abalone.py diff --git a/tests/data/xgboost_abalone/abalone.py b/tests/data/xgboost_abalone/abalone.py new file mode 100644 index 0000000000..a034a80a66 --- /dev/null +++ b/tests/data/xgboost_abalone/abalone.py @@ -0,0 +1,46 @@ +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 diff --git a/tests/integ/test_xgboost.py b/tests/integ/test_xgboost.py index 56fdffab1b..a05b85dad5 100644 --- a/tests/integ/test_xgboost.py +++ b/tests/integ/test_xgboost.py @@ -14,6 +14,7 @@ import os import pytest +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 +49,37 @@ 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" + ) + xgboost.fit( + inputs={"train": train_input}, + job_name=unique_name_from_base(base_job_name) + ) + assert sagemaker_session.sagemaker_client.describe_training_job(TrainingJobName=job_name)[ + "EnableNetworkIsolation" + ] From 7c880a55fd665c99684b07a3b255133dddea9f23 Mon Sep 17 00:00:00 2001 From: Brent Millare Date: Fri, 8 Oct 2021 20:42:48 +0000 Subject: [PATCH 4/7] fix import --- tests/integ/test_xgboost.py | 1 + 1 file changed, 1 insertion(+) diff --git a/tests/integ/test_xgboost.py b/tests/integ/test_xgboost.py index a05b85dad5..50a29cb3c5 100644 --- a/tests/integ/test_xgboost.py +++ b/tests/integ/test_xgboost.py @@ -14,6 +14,7 @@ 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 86f6b08d064820c61d4da7f6840ddf8a207cf1fe Mon Sep 17 00:00:00 2001 From: Brent Millare Date: Fri, 8 Oct 2021 21:03:16 +0000 Subject: [PATCH 5/7] fix job_name name --- tests/integ/test_xgboost.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/tests/integ/test_xgboost.py b/tests/integ/test_xgboost.py index 50a29cb3c5..dffbcbfd08 100644 --- a/tests/integ/test_xgboost.py +++ b/tests/integ/test_xgboost.py @@ -77,9 +77,10 @@ def test_training_with_network_isolation( 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=unique_name_from_base(base_job_name) + job_name=job_name ) assert sagemaker_session.sagemaker_client.describe_training_job(TrainingJobName=job_name)[ "EnableNetworkIsolation" From d533fa7c72791caf04ffe0b002f7182a439d1cbb Mon Sep 17 00:00:00 2001 From: Brent Millare Date: Tue, 19 Oct 2021 02:03:49 +0000 Subject: [PATCH 6/7] fix black-check --- tests/data/xgboost_abalone/abalone.py | 40 +++++++++++++++------------ tests/integ/test_xgboost.py | 7 ++--- 2 files changed, 24 insertions(+), 23 deletions(-) diff --git a/tests/data/xgboost_abalone/abalone.py b/tests/data/xgboost_abalone/abalone.py index a034a80a66..0ba5213c1d 100644 --- a/tests/data/xgboost_abalone/abalone.py +++ b/tests/data/xgboost_abalone/abalone.py @@ -5,35 +5,39 @@ import xgboost as xgb -model_filename = 'xgboost-model' +model_filename = "xgboost-model" -if __name__ == '__main__': +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')) + 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') + 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', + "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) + booster = xgb.train(params=params, dtrain=dtrain, num_boost_round=50) + booster.save_model(args.model_dir + "/" + model_filename) def model_fn(model_dir): diff --git a/tests/integ/test_xgboost.py b/tests/integ/test_xgboost.py index dffbcbfd08..088a55d7f3 100644 --- a/tests/integ/test_xgboost.py +++ b/tests/integ/test_xgboost.py @@ -75,13 +75,10 @@ def test_training_with_network_isolation( train_input = xgboost.sagemaker_session.upload_data( path=os.path.join(DATA_DIR, "xgboost_abalone", "abalone"), - key_prefix="integ-test-data/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 - ) + xgboost.fit(inputs={"train": train_input}, job_name=job_name) assert sagemaker_session.sagemaker_client.describe_training_job(TrainingJobName=job_name)[ "EnableNetworkIsolation" ] From 97d63cd1e4766f8d9657a5b20db7d573f2d606a3 Mon Sep 17 00:00:00 2001 From: Brent Millare Date: Tue, 19 Oct 2021 20:27:58 +0000 Subject: [PATCH 7/7] Revert local test setup changes --- tox.ini | 1 - 1 file changed, 1 deletion(-) diff --git a/tox.ini b/tox.ini index 57d0d88b00..d1a503be88 100644 --- a/tox.ini +++ b/tox.ini @@ -58,7 +58,6 @@ markers = timeout: mark a test as a timeout. [testenv] -basepython = /usr/bin/python3 passenv = AWS_ACCESS_KEY_ID AWS_SECRET_ACCESS_KEY