[Core] Out of Disk prevention

BUILD.bazel

@@ -1431,6 +1431,7 @@ cc_test(
     copts = COPTS,
     tags = ["team:core"],
     deps = [
+        ":ray_common",
         ":ray_util",
         "@com_google_googletest//:gtest_main",
     ],

python/ray/_private/serialization.py

@@ -25,6 +25,7 @@
     ObjectReconstructionFailedError,
     ObjectReconstructionFailedLineageEvictedError,
     ObjectReconstructionFailedMaxAttemptsExceededError,
+    OutOfDiskError,
     OwnerDiedError,
     PlasmaObjectNotAvailable,
     RayActorError,
@@ -276,6 +277,10 @@ def _deserialize_object(self, data, metadata, object_ref):
                 return ObjectFetchTimedOutError(
                     object_ref.hex(), object_ref.owner_address(), object_ref.call_site()
                 )
+            elif error_type == ErrorType.Value("OUT_OF_DISK_ERROR"):
+                return OutOfDiskError(
+                    object_ref.hex(), object_ref.owner_address(), object_ref.call_site()
+                )
             elif error_type == ErrorType.Value("OBJECT_DELETED"):
                 return ReferenceCountingAssertionError(
                     object_ref.hex(), object_ref.owner_address(), object_ref.call_site()

python/ray/_raylet.pyx

@@ -120,6 +120,7 @@ from ray.exceptions import (
     RaySystemError,
     RayTaskError,
     ObjectStoreFullError,
+    OutOfDiskError,
     GetTimeoutError,
     TaskCancelledError,
     AsyncioActorExit,
@@ -166,6 +167,8 @@ cdef int check_status(const CRayStatus& status) nogil except -1:
 
     if status.IsObjectStoreFull():
         raise ObjectStoreFullError(message)
+    elif status.IsOutOfDisk():
+        raise OutOfDiskError(message)
     elif status.IsInterrupted():
         raise KeyboardInterrupt()
     elif status.IsTimedOut():

python/ray/exceptions.py

@@ -341,6 +341,26 @@ def __str__(self):
         )
 
 
+@PublicAPI
+class OutOfDiskError(RayError):
+    """Indicates that the local disk is full.
+
+    This is raised if the attempt to store the object fails
+    because both the object store and disk are full.
+    """
+
+    def __str__(self):
+        # TODO(scv119): expose more disk usage information and link to a doc.
+        return super(OutOfDiskError, self).__str__() + (
+            "\n"
+            "The object cannot be created because the local object store"
+            " is full and the local disk's utilization is over capacity"
+            " (95% by default)."
+            "Tip: Use `df` on this node to check disk usage and "
+            "`ray memory` to check object store memory usage."
+        )
+
+
 @PublicAPI
 class ObjectLostError(RayError):
     """Indicates that the object is lost from distributed memory, due to

python/ray/includes/common.pxd

@@ -108,6 +108,7 @@ cdef extern from "ray/common/status.h" namespace "ray" nogil:
         c_bool IsUnknownError()
         c_bool IsNotImplemented()
         c_bool IsObjectStoreFull()
+        c_bool IsOutOfDisk()
         c_bool IsRedisError()
         c_bool IsTimedOut()
         c_bool IsInterrupted()

python/ray/tests/BUILD

@@ -252,6 +252,7 @@ py_test_module_list(
   files = [
     "test_basic_3.py",
     "test_output.py",
+    "test_out_of_disk_space.py",
     "test_failure_4.py",
     "test_object_spilling.py",
     "test_object_spilling_no_asan.py",

python/ray/tests/test_out_of_disk_space.py

@@ -0,0 +1,221 @@
+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.cluster_utils import Cluster
+
+
+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} ")
+    print(f"current usage: {1.0 - 1.0 * usage.free  / usage.total}")
+    return 1.0 - 1.0 * usage.free / usage.total
+
+
+@contextmanager
+def create_tmp_file(bytes):
+    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(bytes))
+    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):
+        assert get_current_usage() > local_fs_capacity_threshold
+        time.sleep(1)
+        with pytest.raises(ray.exceptions.OutOfDiskError):
+            ray.put(np.random.rand(20 * 1024 * 1024))
+        # delete tmp file to reclaim space back.
+
+    assert get_current_usage() < local_fs_capacity_threshold
+    time.sleep(1)
+    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(shutdown_only):
+    cluster = 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(shutdown_only):
+    cluster = 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)
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main(["-sv", __file__]))

python/ray/tests/test_plasma_unlimited.py

@@ -193,6 +193,8 @@ def test_fallback_allocation_failure(shutdown_only):
         _temp_dir="/dev/shm",
         _system_config={
             "object_spilling_config": json.dumps(file_system_config),
+            # set local fs capacity to 100% so it never errors with out of disk.
+            "local_fs_capacity_threshold": 1,
         },
     )
     shm_size = shutil.disk_usage("/dev/shm").total