feature: modify RLEstimator to use newly generated Ray image (1.6.0) (#2717)

* Modify RLEstimator to use newly generated Ray image (1.6.0)
Changes have been tested locally with Cartpole env.

Rev-1: fix the image tag name and unit tests.

* Fix integration tests due to new Ray version ('webui_host' has been deprecated)

Co-authored-by: Longyu Zhao <[email protected]>
Co-authored-by: Navin Soni <[email protected]>

Author: 3 people

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,

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/unit/sagemaker/image_uris/test_rl.py

@@ -86,7 +86,15 @@ def test_ray_tf(ray_tensorflow_version):
 
 
 def _expected_ray_tf_uri(ray_tf_version, processor):
-    if Version(ray_tf_version) > Version("0.6.5"):
+    if Version(ray_tf_version) > Version("1.0.0"):
+        return expected_uris.framework_uri(
+            "sagemaker-rl-ray-container",
+            _version_for_tag("ray", ray_tf_version, "tf", True),
+            RL_ACCOUNT,
+            py_version="py37",
+            processor=processor,
+        )
+    elif Version(ray_tf_version) > Version("0.6.5"):
         return expected_uris.framework_uri(
             "sagemaker-rl-ray-container",
             _version_for_tag("ray", ray_tf_version, "tf", True),