chore: import upstream snapshot with attribution

python/ray/tests/test_out_of_disk_space.py

@@ -0,0 +1,267 @@
+import os
+import platform
+import shutil
+import sys
+import tempfile
+import time
+from contextlib import contextmanager
+
+import numpy as np
+import pytest
+
+import ray
+from ray._common.test_utils import wait_for_condition
+from ray.util.state import list_cluster_events
+
+
+def calculate_capacity_threshold(disk_capacity_in_bytes):
+    usage = shutil.disk_usage("/tmp")
+    threshold = min(1, 1.0 - 1.0 * (usage.free - disk_capacity_in_bytes) / usage.total)
+    return threshold
+
+
+def get_current_usage():
+    usage = shutil.disk_usage("/tmp")
+    print(f"free: {usage.free} ")
+    usage_percentage = 1.0 - 1.0 * usage.free / usage.total
+    print(f"current usage: {usage_percentage}")
+    return usage_percentage
+
+
+@contextmanager
+def create_tmp_file(n):
+    tmp_dir = tempfile.mkdtemp(dir="/tmp")
+    tmp_path = os.path.join(tmp_dir, "test.txt")
+    with open(tmp_path, "wb") as f:
+        f.write(os.urandom(n))
+    try:
+        yield tmp_path
+    finally:
+        os.remove(tmp_path)
+
+
+@pytest.mark.skipif(platform.system() == "Windows", reason="Not targeting Windows")
+def test_put_out_of_disk(shutdown_only):
+    local_fs_capacity_threshold = calculate_capacity_threshold(200 * 1024 * 1024)
+    ray.init(
+        num_cpus=1,
+        object_store_memory=80 * 1024 * 1024,
+        _system_config={
+            "local_fs_capacity_threshold": local_fs_capacity_threshold,
+            "local_fs_monitor_interval_ms": 10,
+        },
+    )
+    assert get_current_usage() < local_fs_capacity_threshold
+    ref = ray.put(np.random.rand(20 * 1024 * 1024))
+    del ref
+    # create a temp file so that the disk size is over the threshold.
+    # ray.put doesn't work is that fallback allocation uses mmaped file
+    # that doesn't neccssary allocate disk spaces.
+    with create_tmp_file(250 * 1024 * 1024):
+        wait_for_condition(lambda: get_current_usage() > local_fs_capacity_threshold)
+        with pytest.raises(ray.exceptions.OutOfDiskError):
+            ray.put(np.random.rand(20 * 1024 * 1024))
+        # delete tmp file to reclaim space back.
+
+    # Ideally get_current_usage() should immediately reflect the latest disk usage
+    # after the tmp file is deleted, but somehow there is some delays on CI machines
+    # so I use wait_for_condition here.
+    wait_for_condition(lambda: get_current_usage() < local_fs_capacity_threshold)
+    ray.put(np.random.rand(20 * 1024 * 1024))
+
+
+@pytest.mark.skipif(platform.system() == "Windows", reason="Not targeting Windows")
+def test_task_returns(shutdown_only):
+    local_fs_capacity_threshold = calculate_capacity_threshold(10 * 1024 * 1024)
+    ray.init(
+        num_cpus=1,
+        object_store_memory=80 * 1024 * 1024,
+        _system_config={
+            "local_fs_capacity_threshold": local_fs_capacity_threshold,
+            "local_fs_monitor_interval_ms": 10,
+        },
+    )
+
+    # create a temp file so that the disk size is over the threshold.
+    # ray.put doesn't work is that fallback allocation uses mmaped file
+    # that doesn't neccssary allocate disk spaces.
+    with create_tmp_file(250 * 1024 * 1024):
+        assert get_current_usage() > local_fs_capacity_threshold
+        time.sleep(1)
+
+        @ray.remote
+        def foo():
+            time.sleep(1)
+            return np.random.rand(20 * 1024 * 1024)  # 160 MB data
+
+        try:
+            ray.get(foo.remote())
+        except ray.exceptions.RayTaskError as e:
+            assert isinstance(e.cause, ray.exceptions.OutOfDiskError)
+
+
+@pytest.mark.skipif(platform.system() == "Windows", reason="Not targeting Windows")
+def test_task_put(shutdown_only):
+    local_fs_capacity_threshold = calculate_capacity_threshold(1 * 1024 * 1024)
+    ray.init(
+        num_cpus=1,
+        object_store_memory=80 * 1024 * 1024,
+        _system_config={
+            "local_fs_capacity_threshold": local_fs_capacity_threshold,
+            "local_fs_monitor_interval_ms": 10,
+        },
+    )
+
+    # create a temp file so that the disk size is over the threshold.
+    # ray.put doesn't work is that fallback allocation uses mmaped file
+    # that doesn't neccssary allocate disk spaces.
+    with create_tmp_file(250 * 1024 * 1024):
+        assert get_current_usage() > local_fs_capacity_threshold
+        time.sleep(1)
+
+        @ray.remote
+        def foo():
+            ref = ray.put(np.random.rand(20 * 1024 * 1024))  # 160 MB data
+            return ref
+
+        try:
+            ray.get(foo.remote())
+        except ray.exceptions.RayTaskError as e:
+            assert isinstance(e.cause, ray.exceptions.OutOfDiskError)
+
+
+@pytest.mark.skipif(platform.system() == "Windows", reason="Not targeting Windows")
+def test_task_args(ray_start_cluster):
+    cluster = ray_start_cluster
+    cluster.add_node(
+        num_cpus=1,
+        object_store_memory=80 * 1024 * 1024,
+        _system_config={
+            "local_fs_capacity_threshold": 0,
+        },
+        resources={"out_of_memory": 1},
+    )
+    cluster.add_node(
+        num_cpus=1,
+        object_store_memory=200 * 1024 * 1024,
+        resources={"sufficient_memory": 1},
+    )
+    cluster.wait_for_nodes()
+    ray.init(address=cluster.address)
+
+    @ray.remote
+    def foo():
+        return np.random.rand(20 * 1024 * 1024)  # 160 MB data
+
+    @ray.remote
+    def bar(obj):
+        print(obj)
+
+    ref = foo.options(resources={"sufficient_memory": 1}).remote()
+    try:
+        ray.get(bar.options(resources={"out_of_memory": 1}).remote(ref))
+    except ray.exceptions.RayTaskError as e:
+        assert isinstance(e.cause, ray.exceptions.OutOfDiskError)
+
+
+@pytest.mark.skipif(platform.system() == "Windows", reason="Not targeting Windows")
+def test_actor(ray_start_cluster):
+    cluster = ray_start_cluster
+    cluster.add_node(
+        num_cpus=1,
+        object_store_memory=80 * 1024 * 1024,
+        _system_config={
+            "local_fs_capacity_threshold": 0,
+        },
+        resources={"out_of_memory": 1},
+    )
+    cluster.add_node(
+        num_cpus=1,
+        object_store_memory=200 * 1024 * 1024,
+        resources={"sufficient_memory": 1},
+    )
+    cluster.wait_for_nodes()
+    ray.init(address=cluster.address)
+
+    @ray.remote
+    def foo():
+        return np.random.rand(20 * 1024 * 1024)  # 160 MB data
+
+    @ray.remote
+    class Actor:
+        def __init__(self, obj):
+            self._obj = obj
+
+        def foo(self):
+            print(self._obj)
+
+        def args_ood(self, obj):
+            print(obj)
+
+        def return_ood(self):
+            return np.random.rand(20 * 1024 * 1024)
+
+    ref = foo.options(resources={"sufficient_memory": 1}).remote()
+    with pytest.raises(ray.exceptions.RayActorError):
+        a = Actor.options(resources={"out_of_memory": 0.001}).remote(ref)
+        ray.get(a.foo.remote())
+
+    a = Actor.options(resources={"out_of_memory": 1}).remote(1)
+    ray.get(a.foo.remote())
+    try:
+        ray.get(a.args_ood.remote(ref))
+    except ray.exceptions.RayTaskError as e:
+        assert isinstance(e.cause, ray.exceptions.OutOfDiskError)
+
+    ray.get(a.foo.remote())
+    try:
+        ray.get(a.return_ood.remote())
+    except ray.exceptions.RayTaskError as e:
+        assert isinstance(e.cause, ray.exceptions.OutOfDiskError)
+
+
+@pytest.mark.skipif(platform.system() == "Windows", reason="Not targeting Windows")
+def test_ood_events(shutdown_only):
+    local_fs_capacity_threshold = calculate_capacity_threshold(200 * 1024 * 1024)
+    ray.init(
+        num_cpus=1,
+        object_store_memory=80 * 1024 * 1024,
+        _system_config={
+            "local_fs_capacity_threshold": local_fs_capacity_threshold,
+            "local_fs_monitor_interval_ms": 10,
+        },
+    )
+
+    # create a temp file so that the disk size is over the threshold.
+    # ray.put doesn't work is that fallback allocation uses mmaped file
+    # that doesn't neccssary allocate disk spaces.
+    with create_tmp_file(250 * 1024 * 1024):
+        assert get_current_usage() > local_fs_capacity_threshold
+        time.sleep(1)
+
+        @ray.remote
+        def foo():
+            ref = ray.put(np.random.rand(20 * 1024 * 1024))  # 160 MB data
+            return ref
+
+        try:
+            ray.get(foo.remote())
+        except ray.exceptions.RayTaskError as e:
+            assert isinstance(e.cause, ray.exceptions.OutOfDiskError)
+
+    # Give it some time for events to appear.
+    # TODO(core-team): provide some way to wait for events to be flushed.
+    time.sleep(2)
+
+    events = list_cluster_events()
+    print(events)
+    # There could be more than 1 event depending on the test timing.
+    assert len(events) >= 1
+    assert (
+        "Object creation will fail if spilling is required" in events[0]["message"]
+    )  # noqa
+    assert events[0]["severity"] == "ERROR"
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main(["-sv", __file__]))