Ray, Tune, and RLlib support for memory, object_store_memory options (#5226)

python/ray/tests/cluster_utils.py

@@ -62,7 +62,7 @@ class Cluster(object):
         All nodes are by default started with the following settings:
             cleanup=True,
             num_cpus=1,
-            object_store_memory=100 * (2**20) # 100 MB
+            object_store_memory=150 * 1024 * 1024  # 150 MiB
 
         Args:
             node_args: Keyword arguments used in `start_ray_head` and
@@ -74,7 +74,7 @@ class Cluster(object):
         default_kwargs = {
             "num_cpus": 1,
             "num_gpus": 0,
-            "object_store_memory": 100 * (2**20),  # 100 MB
+            "object_store_memory": 150 * 1024 * 1024,  # 150 MiB
         }
         ray_params = ray.parameter.RayParams(**node_args)
         ray_params.update_if_absent(**default_kwargs)

python/ray/tests/conftest.py

@@ -38,7 +38,7 @@ def get_default_fixture_ray_kwargs():
     internal_config = get_default_fixure_internal_config()
     ray_kwargs = {
         "num_cpus": 1,
-        "object_store_memory": 10**8,
+        "object_store_memory": 150 * 1024 * 1024,
         "_internal_config": internal_config,
     }
     return ray_kwargs

python/ray/tests/perf_integration_tests/test_perf_integration.py

@@ -37,7 +37,9 @@ def warmup():
 def test_task_submission(benchmark, num_tasks):
     num_cpus = 16
     ray.init(
-        num_cpus=num_cpus, object_store_memory=10**7, ignore_reinit_error=True)
+        num_cpus=num_cpus,
+        object_store_memory=150 * 1024 * 1024,
+        ignore_reinit_error=True)
     # warm up the plasma store
     warmup()
     benchmark(benchmark_task_submission, num_tasks)
@@ -57,11 +59,11 @@ def test_task_forward(benchmark, num_tasks):
             do_init=True,
             num_nodes=1,
             num_cpus=16,
-            object_store_memory=10**7,
+            object_store_memory=150 * 1024 * 1024,
     ) as cluster:
         cluster.add_node(
             num_cpus=16,
-            object_store_memory=10**7,
+            object_store_memory=150 * 1024 * 1024,
             resources={"my_resource": 100},
         )
 

python/ray/tests/test_actor.py

@@ -444,7 +444,8 @@ def test_actor_deletion(ray_start_regular):
 
 
 def test_actor_deletion_with_gpus(shutdown_only):
-    ray.init(num_cpus=1, num_gpus=1, object_store_memory=int(10**8))
+    ray.init(
+        num_cpus=1, num_gpus=1, object_store_memory=int(150 * 1024 * 1024))
 
     # When an actor that uses a GPU exits, make sure that the GPU resources
     # are released.
@@ -516,7 +517,7 @@ def test_resource_assignment(shutdown_only):
         num_cpus=16,
         num_gpus=1,
         resources={"Custom": 1},
-        object_store_memory=int(10**8))
+        object_store_memory=int(150 * 1024 * 1024))
 
     class Actor(object):
         def __init__(self):
@@ -1296,7 +1297,8 @@ def test_actors_and_tasks_with_gpus(ray_start_cluster):
 def test_actors_and_tasks_with_gpus_version_two(shutdown_only):
     # Create tasks and actors that both use GPUs and make sure that they
     # are given different GPUs
-    ray.init(num_cpus=10, num_gpus=10, object_store_memory=int(10**8))
+    ray.init(
+        num_cpus=10, num_gpus=10, object_store_memory=int(150 * 1024 * 1024))
 
     @ray.remote(num_gpus=1)
     def f():
@@ -1330,7 +1332,8 @@ def test_actors_and_tasks_with_gpus_version_two(shutdown_only):
 
 
 def test_blocking_actor_task(shutdown_only):
-    ray.init(num_cpus=1, num_gpus=1, object_store_memory=int(10**8))
+    ray.init(
+        num_cpus=1, num_gpus=1, object_store_memory=int(150 * 1024 * 1024))
 
     @ray.remote(num_gpus=1)
     def f():
@@ -1740,7 +1743,7 @@ def test_nondeterministic_reconstruction_concurrent_forks(
 
 @pytest.fixture
 def setup_queue_actor():
-    ray.init(num_cpus=1, object_store_memory=int(10**8))
+    ray.init(num_cpus=1, object_store_memory=int(150 * 1024 * 1024))
 
     @ray.remote
     class Queue(object):
@@ -2105,7 +2108,7 @@ def test_creating_more_actors_than_resources(shutdown_only):
 
 
 @pytest.mark.parametrize(
-    "ray_start_object_store_memory", [10**8], indirect=True)
+    "ray_start_object_store_memory", [150 * 1024 * 1024], indirect=True)
 def test_actor_eviction(ray_start_object_store_memory):
     object_store_memory = ray_start_object_store_memory
 

python/ray/tests/test_basic.py

@@ -967,11 +967,9 @@ def test_many_fractional_resources(shutdown_only):
     stop_time = time.time() + 10
     correct_available_resources = False
     while time.time() < stop_time:
-        if ray.available_resources() == {
-                "CPU": 2.0,
-                "GPU": 2.0,
-                "Custom": 2.0,
-        }:
+        if (ray.available_resources()["CPU"] == 2.0
+                and ray.available_resources()["GPU"] == 2.0
+                and ray.available_resources()["Custom"] == 2.0):
             correct_available_resources = True
             break
     if not correct_available_resources:
@@ -2324,6 +2322,9 @@ def test_zero_capacity_deletion_semantics(shutdown_only):
         MAX_RETRY_ATTEMPTS = 5
         retry_count = 0
 
+        del resources["memory"]
+        del resources["object_store_memory"]
+
         while resources and retry_count < MAX_RETRY_ATTEMPTS:
             time.sleep(0.1)
             resources = ray.available_resources()
@@ -2537,8 +2538,9 @@ def test_global_state_api(shutdown_only):
 
     ray.init(num_cpus=5, num_gpus=3, resources={"CustomResource": 1})
 
-    resources = {"CPU": 5, "GPU": 3, "CustomResource": 1}
-    assert ray.cluster_resources() == resources
+    assert ray.cluster_resources()["CPU"] == 5
+    assert ray.cluster_resources()["GPU"] == 3
+    assert ray.cluster_resources()["CustomResource"] == 1
 
     assert ray.objects() == {}
 
@@ -2807,7 +2809,7 @@ def test_initialized_local_mode(shutdown_only_with_initialization_check):
 
 
 def test_wait_reconstruction(shutdown_only):
-    ray.init(num_cpus=1, object_store_memory=10**8)
+    ray.init(num_cpus=1, object_store_memory=int(10**8))
 
     @ray.remote
     def f():
@@ -3025,7 +3027,7 @@ def test_shutdown_disconnect_global_state():
 
 
 @pytest.mark.parametrize(
-    "ray_start_object_store_memory", [10**8], indirect=True)
+    "ray_start_object_store_memory", [150 * 1024 * 1024], indirect=True)
 def test_redis_lru_with_set(ray_start_object_store_memory):
     x = np.zeros(8 * 10**7, dtype=np.uint8)
     x_id = ray.put(x)

python/ray/tests/test_cython.py

@@ -16,7 +16,7 @@ def get_ray_result(cython_func, *args):
 
 class CythonTest(unittest.TestCase):
     def setUp(self):
-        ray.init(object_store_memory=int(10**8))
+        ray.init(object_store_memory=int(150 * 1024 * 1024))
 
     def tearDown(self):
         ray.shutdown()

python/ray/tests/test_failure.py

@@ -725,7 +725,7 @@ def test_connect_with_disconnected_node(shutdown_only):
 @pytest.mark.parametrize(
     "ray_start_cluster_head", [{
         "num_cpus": 5,
-        "object_store_memory": 10**7
+        "object_store_memory": 10**8
     }],
     indirect=True)
 @pytest.mark.parametrize("num_actors", [1, 2, 5])
@@ -733,7 +733,7 @@ def test_parallel_actor_fill_plasma_retry(ray_start_cluster_head, num_actors):
     @ray.remote
     class LargeMemoryActor(object):
         def some_expensive_task(self):
-            return np.zeros(10**7 // 2, dtype=np.uint8)
+            return np.zeros(10**8 // 2, dtype=np.uint8)
 
     actors = [LargeMemoryActor.remote() for _ in range(num_actors)]
     for _ in range(10):
@@ -745,14 +745,14 @@ def test_parallel_actor_fill_plasma_retry(ray_start_cluster_head, num_actors):
 @pytest.mark.parametrize(
     "ray_start_cluster_head", [{
         "num_cpus": 2,
-        "object_store_memory": 10**7
+        "object_store_memory": 10**8
     }],
     indirect=True)
 def test_fill_plasma_exception(ray_start_cluster_head):
     @ray.remote
     class LargeMemoryActor(object):
         def some_expensive_task(self):
-            return np.zeros(10**7 + 2, dtype=np.uint8)
+            return np.zeros(10**8 + 2, dtype=np.uint8)
 
         def test(self):
             return 1
@@ -764,4 +764,4 @@ def test_fill_plasma_exception(ray_start_cluster_head):
     ray.get(actor.test.remote())
 
     with pytest.raises(plasma.PlasmaStoreFull):
-        ray.put(np.zeros(10**7 + 2, dtype=np.uint8))
+        ray.put(np.zeros(10**8 + 2, dtype=np.uint8))

python/ray/tests/test_memory_limits.py

@@ -0,0 +1,85 @@
+import numpy as np
+import unittest
+
+import ray
+import pyarrow
+
+MB = 1024 * 1024
+
+OBJECT_EVICTED = ray.exceptions.UnreconstructableError
+OBJECT_TOO_LARGE = pyarrow._plasma.PlasmaStoreFull
+
+
+@ray.remote
+class LightActor(object):
+    def __init__(self):
+        pass
+
+    def sample(self):
+        return "tiny_return_value"
+
+
+@ray.remote
+class GreedyActor(object):
+    def __init__(self):
+        pass
+
+    def sample(self):
+        return np.zeros(20 * MB, dtype=np.uint8)
+
+
+class TestMemoryLimits(unittest.TestCase):
+    def testWithoutQuota(self):
+        self.assertRaises(OBJECT_EVICTED, lambda: self._run(None, None, None))
+        self.assertRaises(OBJECT_EVICTED,
+                          lambda: self._run(100 * MB, None, None))
+        self.assertRaises(OBJECT_EVICTED,
+                          lambda: self._run(None, 100 * MB, None))
+
+    def testQuotasProtectSelf(self):
+        self._run(100 * MB, 100 * MB, None)
+
+    def testQuotasProtectOthers(self):
+        self._run(None, None, 100 * MB)
+
+    def testQuotaTooLarge(self):
+        self.assertRaisesRegexp(ray.memory_monitor.RayOutOfMemoryError,
+                                ".*Failed to set object_store_memory.*",
+                                lambda: self._run(300 * MB, None, None))
+
+    def testTooLargeAllocation(self):
+        try:
+            ray.init(num_cpus=1, driver_object_store_memory=100 * MB)
+            ray.put(np.zeros(50 * MB, dtype=np.uint8))
+            self.assertRaises(
+                OBJECT_TOO_LARGE,
+                lambda: ray.put(np.zeros(200 * MB, dtype=np.uint8)))
+        finally:
+            ray.shutdown()
+
+    def _run(self, driver_quota, a_quota, b_quota):
+        print("*** Testing ***", driver_quota, a_quota, b_quota)
+        try:
+            ray.init(
+                num_cpus=1,
+                object_store_memory=300 * MB,
+                driver_object_store_memory=driver_quota)
+            z = ray.put("hi")
+            a = LightActor._remote(object_store_memory=a_quota)
+            b = GreedyActor._remote(object_store_memory=b_quota)
+            for _ in range(5):
+                r_a = a.sample.remote()
+                for _ in range(20):
+                    ray.get(b.sample.remote())
+                ray.get(r_a)
+            ray.get(z)
+        except Exception as e:
+            print("Raised exception", type(e), e)
+            raise e
+        finally:
+            print(ray.worker.global_worker.plasma_client.debug_string())
+            ray.shutdown()
+
+
+if __name__ == "__main__":
+    unittest.main(verbosity=2)

python/ray/tests/test_memory_scheduling.py

@@ -0,0 +1,155 @@
+import numpy as np
+import unittest
+
+import ray
+from ray import tune
+from ray.rllib import _register_all
+
+MB = 1024 * 1024
+
+
+@ray.remote(memory=100 * MB)
+class Actor(object):
+    def __init__(self):
+        pass
+
+    def ping(self):
+        return "ok"
+
+
+@ray.remote(object_store_memory=100 * MB)
+class Actor2(object):
+    def __init__(self):
+        pass
+
+    def ping(self):
+        return "ok"
+
+
+def train_oom(config, reporter):
+    ray.put(np.zeros(200 * 1024 * 1024))
+    reporter(result=123)
+
+
+class TestMemoryScheduling(unittest.TestCase):
+    def testMemoryRequest(self):
+        try:
+            ray.init(num_cpus=1, memory=200 * MB)
+            # fits first 2
+            a = Actor.remote()
+            b = Actor.remote()
+            ok, _ = ray.wait(
+                [a.ping.remote(), b.ping.remote()],
+                timeout=60.0,
+                num_returns=2)
+            self.assertEqual(len(ok), 2)
+            # does not fit
+            c = Actor.remote()
+            ok, _ = ray.wait([c.ping.remote()], timeout=5.0)
+            self.assertEqual(len(ok), 0)
+        finally:
+            ray.shutdown()
+
+    def testObjectStoreMemoryRequest(self):
+        try:
+            ray.init(num_cpus=1, object_store_memory=300 * MB)
+            # fits first 2 (70% allowed)
+            a = Actor2.remote()
+            b = Actor2.remote()
+            ok, _ = ray.wait(
+                [a.ping.remote(), b.ping.remote()],
+                timeout=60.0,
+                num_returns=2)
+            self.assertEqual(len(ok), 2)
+            # does not fit
+            c = Actor2.remote()
+            ok, _ = ray.wait([c.ping.remote()], timeout=5.0)
+            self.assertEqual(len(ok), 0)
+        finally:
+            ray.shutdown()
+
+    def testTuneDriverHeapLimit(self):
+        try:
+            _register_all()
+            result = tune.run(
+                "PG",
+                stop={"timesteps_total": 10000},
+                config={
+                    "env": "CartPole-v0",
+                    "memory": 100 * 1024 * 1024,  # too little
+                },
+                raise_on_failed_trial=False)
+            self.assertEqual(result.trials[0].status, "ERROR")
+            self.assertTrue(
+                "RayOutOfMemoryError: Heap memory usage for ray_PG_" in
+                result.trials[0].error_msg)
+        finally:
+            ray.shutdown()
+
+    def testTuneDriverStoreLimit(self):
+        try:
+            _register_all()
+            self.assertRaisesRegexp(
+                ray.tune.error.TuneError,
+                ".*Insufficient cluster resources.*",
+                lambda: tune.run(
+                    "PG",
+                    stop={"timesteps_total": 10000},
+                    config={
+                        "env": "CartPole-v0",
+                        # too large
+                        "object_store_memory": 10000 * 1024 * 1024,
+                    }))
+        finally:
+            ray.shutdown()
+
+    def testTuneWorkerHeapLimit(self):
+        try:
+            _register_all()
+            result = tune.run(
+                "PG",
+                stop={"timesteps_total": 10000},
+                config={
+                    "env": "CartPole-v0",
+                    "num_workers": 1,
+                    "memory_per_worker": 100 * 1024 * 1024,  # too little
+                },
+                raise_on_failed_trial=False)
+            self.assertEqual(result.trials[0].status, "ERROR")
+            self.assertTrue(
+                "RayOutOfMemoryError: Heap memory usage for ray_Rollout" in
+                result.trials[0].error_msg)
+        finally:
+            ray.shutdown()
+
+    def testTuneWorkerStoreLimit(self):
+        try:
+            _register_all()
+            self.assertRaisesRegexp(
+                ray.tune.error.TuneError,
+                ".*Insufficient cluster resources.*",
+                lambda:
+                tune.run("PG", stop={"timesteps_total": 0}, config={
+                    "env": "CartPole-v0",
+                    "num_workers": 1,
+                    # too large
+                    "object_store_memory_per_worker": 10000 * 1024 * 1024,
+                }))
+        finally:
+            ray.shutdown()
+
+    def testTuneObjectLimitApplied(self):
+        try:
+            result = tune.run(
+                train_oom,
+                resources_per_trial={"object_store_memory": 150 * 1024 * 1024},
+                raise_on_failed_trial=False)
+            self.assertTrue(result.trials[0].status, "ERROR")
+            self.assertTrue("PlasmaStoreFull: object does not fit" in
+                            result.trials[0].error_msg)
+        finally:
+            ray.shutdown()
+
+
+if __name__ == "__main__":
+    unittest.main(verbosity=2)

python/ray/tests/test_multi_node_2.py

@@ -73,6 +73,15 @@ def verify_load_metrics(monitor, expected_resource_usage=None, timeout=10):
         monitor.process_messages()
         resource_usage = monitor.load_metrics.get_resource_usage()
 
+        if "memory" in resource_usage[1]:
+            del resource_usage[1]["memory"]
+        if "object_store_memory" in resource_usage[2]:
+            del resource_usage[1]["object_store_memory"]
+        if "memory" in resource_usage[2]:
+            del resource_usage[2]["memory"]
+        if "object_store_memory" in resource_usage[2]:
+            del resource_usage[2]["object_store_memory"]
+
         if expected_resource_usage is None:
             if all(x for x in resource_usage[1:]):
                 break

python/ray/tests/test_object_manager.py

@@ -52,11 +52,11 @@ def test_object_broadcast(ray_start_cluster_with_resource):
     def f(x):
         return
 
-    x = np.zeros(10**8, dtype=np.uint8)
+    x = np.zeros(150 * 1024 * 1024, dtype=np.uint8)
 
     @ray.remote
     def create_object():
-        return np.zeros(10**8, dtype=np.uint8)
+        return np.zeros(150 * 1024 * 1024, dtype=np.uint8)
 
     object_ids = []
 
@@ -219,7 +219,7 @@ def test_object_transfer_retry(ray_start_cluster):
         "object_manager_pull_timeout_ms": repeated_push_delay * 1000 / 4,
         "object_manager_default_chunk_size": 1000
     })
-    object_store_memory = 10**8
+    object_store_memory = 150 * 1024 * 1024
     cluster.add_node(
         object_store_memory=object_store_memory, _internal_config=config)
     cluster.add_node(

python/ray/tests/test_stress.py

@@ -25,7 +25,7 @@ def ray_start_sharded(request):
 
     # Start the Ray processes.
     ray.init(
-        object_store_memory=int(0.1 * 10**9),
+        object_store_memory=int(0.5 * 10**9),
         num_cpus=10,
         num_redis_shards=num_redis_shards,
         redis_max_memory=10**7)
@@ -200,7 +200,7 @@ def test_wait(ray_start_combination):
 def ray_start_reconstruction(request):
     num_nodes = request.param
 
-    plasma_store_memory = int(0.1 * 10**9)
+    plasma_store_memory = int(0.5 * 10**9)
 
     cluster = Cluster(
         initialize_head=True,

python/ray/tests/test_unreconstructable_errors.py

@@ -10,7 +10,10 @@ import ray
 
 class TestUnreconstructableErrors(unittest.TestCase):
     def setUp(self):
-        ray.init(object_store_memory=10000000, redis_max_memory=10000000)
+        ray.init(
+            num_cpus=1,
+            object_store_memory=150 * 1024 * 1024,
+            redis_max_memory=10000000)
 
     def tearDown(self):
         ray.shutdown()
@@ -18,8 +21,8 @@ class TestUnreconstructableErrors(unittest.TestCase):
     def testDriverPutEvictedCannotReconstruct(self):
         x_id = ray.put(np.zeros(1 * 1024 * 1024))
         ray.get(x_id)
-        for _ in range(10):
-            ray.put(np.zeros(1 * 1024 * 1024))
+        for _ in range(20):
+            ray.put(np.zeros(10 * 1024 * 1024))
         self.assertRaises(ray.exceptions.UnreconstructableError,
                           lambda: ray.get(x_id))