Merge branch 'master' into jsonget-deprecation

Author: navinsoni

CHANGELOG.md

@@ -1,5 +1,18 @@
 # Changelog
 
+## v2.65.0 (2021-10-21)
+
+### Features
+
+ * modify RLEstimator to use newly generated Ray image (1.6.0)
+ * network isolation mode for xgboost
+ * update clarify imageURI for PDT
+
+### Bug Fixes and Other Changes
+
+ * retry downstream_trials test
+ * Add retries to pipeline execution
+
 ## v2.64.0 (2021-10-20)
 
 ### Deprecations and Removals

VERSION

@@ -1 +1 @@
-2.64.1.dev0
+2.65.1.dev0

src/sagemaker/image_uri_config/ray-pytorch.json

@@ -21,6 +21,26 @@
             },
             "repository": "sagemaker-rl-ray-container",
             "tag_prefix": "ray-0.8.5-torch"
+        },
+        "1.6.0": {
+            "py_versions": ["py36"],
+            "registries": {
+                "ap-northeast-1": "462105765813",
+                "ap-northeast-2": "462105765813",
+                "ap-south-1": "462105765813",
+                "ap-southeast-1": "462105765813",
+                "ap-southeast-2": "462105765813",
+                "ca-central-1": "462105765813",
+                "eu-central-1": "462105765813",
+                "eu-west-1": "462105765813",
+                "eu-west-2": "462105765813",
+                "us-east-1": "462105765813",
+                "us-east-2": "462105765813",
+                "us-west-1": "462105765813",
+                "us-west-2": "462105765813"
+            },
+            "repository": "sagemaker-rl-ray-container",
+            "tag_prefix": "ray-1.6.0-torch"
         }
     }
 }

src/sagemaker/image_uri_config/ray-tensorflow.json

@@ -165,6 +165,26 @@
             },
             "repository": "sagemaker-rl-ray-container",
             "tag_prefix": "ray-0.8.5-tf"
+        },
+        "1.6.0": {
+            "py_versions": ["py37"],
+            "registries": {
+                "ap-northeast-1": "462105765813",
+                "ap-northeast-2": "462105765813",
+                "ap-south-1": "462105765813",
+                "ap-southeast-1": "462105765813",
+                "ap-southeast-2": "462105765813",
+                "ca-central-1": "462105765813",
+                "eu-central-1": "462105765813",
+                "eu-west-1": "462105765813",
+                "eu-west-2": "462105765813",
+                "us-east-1": "462105765813",
+                "us-east-2": "462105765813",
+                "us-west-1": "462105765813",
+                "us-west-2": "462105765813"
+            },
+            "repository": "sagemaker-rl-ray-container",
+            "tag_prefix": "ray-1.6.0-tf"
         }
     }
 }

src/sagemaker/rl/estimator.py

@@ -45,6 +45,7 @@
         "0.6": {"tensorflow": "1.12"},
         "0.8.2": {"tensorflow": "2.1"},
         "0.8.5": {"tensorflow": "2.1", "pytorch": "1.5"},
+        "1.6.0": {"tensorflow": "2.5.0", "pytorch": "1.8.1"},
     },
 }
 
@@ -69,7 +70,7 @@ class RLEstimator(Framework):
 
     COACH_LATEST_VERSION_TF = "0.11.1"
     COACH_LATEST_VERSION_MXNET = "0.11.0"
-    RAY_LATEST_VERSION = "0.8.5"
+    RAY_LATEST_VERSION = "1.6.0"
 
     def __init__(
         self,

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/ray_cartpole/train_ray.py

@@ -5,10 +5,10 @@
 from ray.tune.logger import pretty_print
 
 # Based on https://github.com/ray-project/ray/blob/master/doc/source/rllib-training.rst#python-api
-ray.init(log_to_driver=False, webui_host="127.0.0.1")
+ray.init(log_to_driver=False)
 config = ppo.DEFAULT_CONFIG.copy()
 config["num_gpus"] = int(os.environ.get("SM_NUM_GPUS", 0))
-checkpoint_dir = os.environ.get("SM_MODEL_DIR", "/Users/nadzeya/gym")
+checkpoint_dir = os.environ.get("SM_MODEL_DIR", "/tmp")
 config["num_workers"] = 1
 agent = ppo.PPOTrainer(config=config, env="CartPole-v0")
 

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/sagemaker/lineage/helpers.py

@@ -15,6 +15,7 @@
 
 import uuid
 from datetime import datetime
+import time
 
 
 def name():
@@ -30,3 +31,16 @@ def names():
         )
         for i in range(3)
     ]
+
+
+def retry(callable, num_attempts=8):
+    assert num_attempts >= 1
+    for i in range(num_attempts):
+        try:
+            return callable()
+        except Exception as ex:
+            if i == num_attempts - 1:
+                raise ex
+            print("Retrying", ex)
+            time.sleep(2 ** i)
+    assert False, "logic error in retry"

tests/integ/sagemaker/lineage/test_artifact.py

@@ -20,6 +20,7 @@
 import pytest
 
 from sagemaker.lineage import artifact
+from tests.integ.sagemaker.lineage.helpers import retry
 
 
 def test_create_delete(artifact_obj):
@@ -103,14 +104,20 @@ def test_list_by_type(artifact_objs, sagemaker_session):
 
 def test_downstream_trials(trial_associated_artifact, trial_obj, sagemaker_session):
     # allow trial components to index, 30 seconds max
-    for i in range(3):
-        time.sleep(10)
-        trials = trial_associated_artifact.downstream_trials(sagemaker_session=sagemaker_session)
-        if len(trials) > 0:
-            break
+    def validate():
+        for i in range(3):
+            time.sleep(10)
+            trials = trial_associated_artifact.downstream_trials(
+                sagemaker_session=sagemaker_session
+            )
+            logging.info(f"Found {len(trials)} downstream trials.")
+            if len(trials) > 0:
+                break
+
+        assert len(trials) == 1
+        assert trial_obj.trial_name in trials
 
-    assert len(trials) == 1
-    assert trial_obj.trial_name in trials
+    retry(validate, num_attempts=3)
 
 
 @pytest.mark.timeout(30)

tests/integ/test_workflow.py

@@ -18,6 +18,7 @@
 import subprocess
 import time
 import uuid
+import logging
 
 from contextlib import contextmanager
 import pytest
@@ -75,6 +76,7 @@
 from sagemaker.feature_store.feature_group import FeatureGroup, FeatureDefinition, FeatureTypeEnum
 from tests.integ import DATA_DIR
 from tests.integ.kms_utils import get_or_create_kms_key
+from tests.integ.retry import retries
 
 
 def ordered(obj):
@@ -1850,47 +1852,57 @@ def test_training_job_with_debugger_and_profiler(
         sagemaker_session=sagemaker_session,
     )
 
-    try:
-        response = pipeline.create(role)
-        create_arn = response["PipelineArn"]
-
-        execution = pipeline.start()
-        response = execution.describe()
-        assert response["PipelineArn"] == create_arn
-
+    for _ in retries(
+        max_retry_count=5,
+        exception_message_prefix="Waiting for a successful execution of pipeline",
+        seconds_to_sleep=10,
+    ):
         try:
-            execution.wait(delay=10, max_attempts=60)
-        except WaiterError:
-            pass
-        execution_steps = execution.list_steps()
+            response = pipeline.create(role)
+            create_arn = response["PipelineArn"]
 
-        assert len(execution_steps) == 1
-        assert execution_steps[0].get("FailureReason", "") == ""
-        assert execution_steps[0]["StepName"] == "pytorch-train"
-        assert execution_steps[0]["StepStatus"] == "Succeeded"
+            execution = pipeline.start()
+            response = execution.describe()
+            assert response["PipelineArn"] == create_arn
 
-        training_job_arn = execution_steps[0]["Metadata"]["TrainingJob"]["Arn"]
-        job_description = sagemaker_session.sagemaker_client.describe_training_job(
-            TrainingJobName=training_job_arn.split("/")[1]
-        )
+            try:
+                execution.wait(delay=10, max_attempts=60)
+            except WaiterError:
+                pass
+            execution_steps = execution.list_steps()
 
-        for index, rule in enumerate(rules):
-            config = job_description["DebugRuleConfigurations"][index]
-            assert config["RuleConfigurationName"] == rule.name
-            assert config["RuleEvaluatorImage"] == rule.image_uri
-            assert config["VolumeSizeInGB"] == 0
-            assert (
-                config["RuleParameters"]["rule_to_invoke"] == rule.rule_parameters["rule_to_invoke"]
+            assert len(execution_steps) == 1
+            failure_reason = execution_steps[0].get("FailureReason", "")
+            if failure_reason != "":
+                logging.error(f"Pipeline execution failed with error: {failure_reason}.Retrying..")
+                continue
+            assert execution_steps[0]["StepName"] == "pytorch-train"
+            assert execution_steps[0]["StepStatus"] == "Succeeded"
+
+            training_job_arn = execution_steps[0]["Metadata"]["TrainingJob"]["Arn"]
+            job_description = sagemaker_session.sagemaker_client.describe_training_job(
+                TrainingJobName=training_job_arn.split("/")[1]
             )
-        assert job_description["DebugHookConfig"] == debugger_hook_config._to_request_dict()
 
-        assert job_description["ProfilingStatus"] == "Enabled"
-        assert job_description["ProfilerConfig"]["ProfilingIntervalInMilliseconds"] == 500
-    finally:
-        try:
-            pipeline.delete()
-        except Exception:
-            pass
+            for index, rule in enumerate(rules):
+                config = job_description["DebugRuleConfigurations"][index]
+                assert config["RuleConfigurationName"] == rule.name
+                assert config["RuleEvaluatorImage"] == rule.image_uri
+                assert config["VolumeSizeInGB"] == 0
+                assert (
+                    config["RuleParameters"]["rule_to_invoke"]
+                    == rule.rule_parameters["rule_to_invoke"]
+                )
+            assert job_description["DebugHookConfig"] == debugger_hook_config._to_request_dict()
+
+            assert job_description["ProfilingStatus"] == "Enabled"
+            assert job_description["ProfilerConfig"]["ProfilingIntervalInMilliseconds"] == 500
+            break
+        finally:
+            try:
+                pipeline.delete()
+            except Exception:
+                pass
 
 
 def test_two_processing_job_depends_on(

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