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

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)