Add integ tests for xgboost net iso

Author: Brent Millare

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

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"
+        ]