chore: import upstream snapshot with attribution

python/ray/tests/test_memory_scheduling.py

@@ -0,0 +1,80 @@
+import sys
+import time
+
+import numpy as np
+import pytest
+
+import ray
+from ray._common.test_utils import wait_for_condition
+
+MB = 1024 * 1024
+
+
+def object_store_memory(a, delta=MB):
+    object_store_memory = ray.available_resources().get("object_store_memory")
+    if object_store_memory is None:
+        return False
+    b = ray.available_resources()["object_store_memory"]
+    return abs(a - b) < delta
+
+
+@ray.remote(memory=100 * MB)
+class Actor:
+    def __init__(self):
+        pass
+
+    def ping(self):
+        return "ok"
+
+
+def test_memory_request():
+    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
+        )
+        assert len(ok) == 2
+        # does not fit
+        c = Actor.remote()
+        ok, _ = ray.wait([c.ping.remote()], timeout=5.0)
+        assert len(ok) == 0
+    finally:
+        ray.shutdown()
+
+
+def test_object_store_memory_reporting():
+    try:
+        ray.init(num_cpus=1, object_store_memory=500 * MB)
+        wait_for_condition(lambda: object_store_memory(500 * MB))
+        x1 = ray.put(np.zeros(150 * 1024 * 1024, dtype=np.uint8))
+        wait_for_condition(lambda: object_store_memory(350 * MB))
+        x2 = ray.put(np.zeros(75 * 1024 * 1024, dtype=np.uint8))
+        wait_for_condition(lambda: object_store_memory(275 * MB))
+        del x1
+        del x2
+        wait_for_condition(lambda: object_store_memory(500 * MB))
+    finally:
+        ray.shutdown()
+
+
+def test_object_store_memory_reporting_task():
+    @ray.remote
+    def f(x):
+        time.sleep(60)
+
+    try:
+        ray.init(num_cpus=1, object_store_memory=500 * MB)
+        wait_for_condition(lambda: object_store_memory(500 * MB))
+        x1 = f.remote(np.zeros(150 * 1024 * 1024, dtype=np.uint8))
+        wait_for_condition(lambda: object_store_memory(350 * MB))
+        ray.cancel(x1, force=True)
+        wait_for_condition(lambda: object_store_memory(500 * MB))
+    finally:
+        ray.shutdown()
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main(["-sv", __file__]))