Repository: ray-project/ray

Issue:
[Object spilling] Raylet automatically reloads spilled objects back into object store




Discussion:

Resolving patch:
diff --git a/python/ray/experimental/__init__.py b/python/ray/experimental/__init__.py
--- a/python/ray/experimental/__init__.py
+++ b/python/ray/experimental/__init__.py
@@ -1,7 +1,6 @@
 from .dynamic_resources import set_resource
-from .object_spilling import force_spill_objects, force_restore_spilled_objects
+from .object_spilling import force_spill_objects
 __all__ = [
     "set_resource",
     "force_spill_objects",
-    "force_restore_spilled_objects",
 ]
diff --git a/python/ray/experimental/object_spilling.py b/python/ray/experimental/object_spilling.py
--- a/python/ray/experimental/object_spilling.py
+++ b/python/ray/experimental/object_spilling.py
@@ -16,20 +16,3 @@ def force_spill_objects(object_refs):
                 f"Attempting to call `force_spill_objects` on the "
                 f"value {object_ref}, which is not an ray.ObjectRef.")
     return core_worker.force_spill_objects(object_refs)
-
-
-def force_restore_spilled_objects(object_refs):
-    """Force restoring objects from external storage.
-
-    Args:
-        object_refs: Object refs of the objects to be
-            restored.
-    """
-    core_worker = ray.worker.global_worker.core_worker
-    # Make sure that the values are object refs.
-    for object_ref in object_refs:
-        if not isinstance(object_ref, ray.ObjectRef):
-            raise TypeError(
-                f"Attempting to call `force_restore_spilled_objects` on the "
-                f"value {object_ref}, which is not an ray.ObjectRef.")
-    return core_worker.force_restore_spilled_objects(object_refs)


Test patch:
diff --git a/python/ray/tests/test_object_spilling.py b/python/ray/tests/test_object_spilling.py
--- a/python/ray/tests/test_object_spilling.py
+++ b/python/ray/tests/test_object_spilling.py
@@ -1,5 +1,7 @@
 import json
 import random
+import platform
+import sys
 import time
 
 import numpy as np
@@ -7,6 +9,8 @@
 import ray
 
 
+@pytest.mark.skipif(
+    platform.system() == "Windows", reason="Failing on Windows.")
 def test_spill_objects_manually(shutdown_only):
     # Limit our object store to 75 MiB of memory.
     ray.init(
@@ -24,7 +28,6 @@ def test_spill_objects_manually(shutdown_only):
     arr = np.random.rand(1024 * 1024)  # 8 MB data
     replay_buffer = []
     pinned_objects = set()
-    spilled_objects = set()
 
     # Create objects of more than 200 MiB.
     for _ in range(25):
@@ -37,23 +40,21 @@ def test_spill_objects_manually(shutdown_only):
             except ray.exceptions.ObjectStoreFullError:
                 ref_to_spill = pinned_objects.pop()
                 ray.experimental.force_spill_objects([ref_to_spill])
-                spilled_objects.add(ref_to_spill)
 
     # Spill 2 more objects so we will always have enough space for
     # restoring objects back.
     refs_to_spill = (pinned_objects.pop(), pinned_objects.pop())
     ray.experimental.force_spill_objects(refs_to_spill)
-    spilled_objects.update(refs_to_spill)
 
     # randomly sample objects
     for _ in range(100):
         ref = random.choice(replay_buffer)
-        if ref in spilled_objects:
-            ray.experimental.force_restore_spilled_objects([ref])
         sample = ray.get(ref)
         assert np.array_equal(sample, arr)
 
 
+@pytest.mark.skipif(
+    platform.system() == "Windows", reason="Failing on Windows.")
 def test_spill_objects_manually_from_workers(shutdown_only):
     # Limit our object store to 100 MiB of memory.
     ray.init(
@@ -71,15 +72,22 @@ def test_spill_objects_manually_from_workers(shutdown_only):
 
     @ray.remote
     def _worker():
-        arr = np.random.rand(100 * 1024)
+        arr = np.random.rand(1024 * 1024)  # 8 MB data
         ref = ray.put(arr)
         ray.experimental.force_spill_objects([ref])
-        ray.experimental.force_restore_spilled_objects([ref])
-        assert np.array_equal(ray.get(ref), arr)
+        return ref
 
-    ray.get([_worker.remote() for _ in range(50)])
+    # Create objects of more than 200 MiB.
+    replay_buffer = [ray.get(_worker.remote()) for _ in range(25)]
+    values = {ref: np.copy(ray.get(ref)) for ref in replay_buffer}
+    # Randomly sample objects.
+    for _ in range(100):
+        ref = random.choice(replay_buffer)
+        sample = ray.get(ref)
+        assert np.array_equal(sample, values[ref])
 
 
+@pytest.mark.skip(reason="Not implemented yet.")
 def test_spill_objects_manually_with_workers(shutdown_only):
     # Limit our object store to 75 MiB of memory.
     ray.init(
@@ -107,27 +115,29 @@ def _worker(object_refs):
         assert np.array_equal(restored, arr)
 
 
+@pytest.mark.skipif(
+    platform.system() == "Windows", reason="Failing on Windows.")
 @pytest.mark.parametrize(
     "ray_start_cluster_head", [{
         "num_cpus": 0,
         "object_store_memory": 75 * 1024 * 1024,
-        "_object_spilling_config": {
+        "object_spilling_config": {
             "type": "filesystem",
             "params": {
                 "directory_path": "/tmp"
             }
         },
-        "_system_config": json.dumps({
+        "_system_config": {
             "object_store_full_max_retries": 0,
             "max_io_workers": 4,
-        }),
+        },
     }],
     indirect=True)
 def test_spill_remote_object(ray_start_cluster_head):
     cluster = ray_start_cluster_head
     cluster.add_node(
         object_store_memory=75 * 1024 * 1024,
-        _object_spilling_config={
+        object_spilling_config={
             "type": "filesystem",
             "params": {
                 "directory_path": "/tmp"
@@ -138,23 +148,33 @@ def test_spill_remote_object(ray_start_cluster_head):
     def put():
         return np.random.rand(5 * 1024 * 1024)  # 40 MB data
 
-    # Create 2 objects. Only 1 should fit.
+    @ray.remote
+    def depends(arg):
+        return
+
     ref = put.remote()
-    ray.get(ref)
+    copy = np.copy(ray.get(ref))
+    # Evict local copy.
+    ray.put(np.random.rand(5 * 1024 * 1024))  # 40 MB data
+    # Remote copy should not fit.
     with pytest.raises(ray.exceptions.RayTaskError):
         ray.get(put.remote())
-    time.sleep(1)
     # Spill 1 object. The second should now fit.
     ray.experimental.force_spill_objects([ref])
     ray.get(put.remote())
 
-    # TODO(swang): Restoring from the object directory is not yet supported.
-    # ray.experimental.force_restore_spilled_objects([ref])
-    # sample = ray.get(ref)
-    # assert np.array_equal(sample, copy)
+    sample = ray.get(ref)
+    assert np.array_equal(sample, copy)
+    # Evict the spilled object.
+    del sample
+    ray.get(put.remote())
+    ray.put(np.random.rand(5 * 1024 * 1024))  # 40 MB data
+
+    # Test passing the spilled object as an arg to another task.
+    ray.get(depends.remote(ref))
 
 
-@pytest.mark.skip(reason="have not been fully implemented")
+@pytest.mark.skip(reason="Not implemented yet.")
 def test_spill_objects_automatically(shutdown_only):
     # Limit our object store to 75 MiB of memory.
     ray.init(
@@ -185,3 +205,7 @@ def test_spill_objects_automatically(shutdown_only):
         ref = random.choice(replay_buffer)
         sample = ray.get(ref, timeout=0)
         assert np.array_equal(sample, arr)
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main(["-sv", __file__]))
diff --git a/src/ray/core_worker/reference_count_test.cc b/src/ray/core_worker/reference_count_test.cc
--- a/src/ray/core_worker/reference_count_test.cc
+++ b/src/ray/core_worker/reference_count_test.cc
@@ -1987,21 +1987,22 @@ TEST_F(ReferenceCountLineageEnabledTest, TestPlasmaLocation) {
   ObjectID borrowed_id = ObjectID::FromRandom();
   rc->AddLocalReference(borrowed_id, "");
   NodeID pinned_at;
-  ASSERT_FALSE(rc->IsPlasmaObjectPinned(borrowed_id, &pinned_at));
+  bool spilled;
+  ASSERT_FALSE(rc->IsPlasmaObjectPinnedOrSpilled(borrowed_id, &pinned_at, &spilled));
 
   ObjectID id = ObjectID::FromRandom();
   NodeID node_id = NodeID::FromRandom();
   rc->AddOwnedObject(id, {}, rpc::Address(), "", 0, true);
   rc->AddLocalReference(id, "");
   ASSERT_TRUE(rc->SetDeleteCallback(id, callback));
-  ASSERT_TRUE(rc->IsPlasmaObjectPinned(id, &pinned_at));
+  ASSERT_TRUE(rc->IsPlasmaObjectPinnedOrSpilled(id, &pinned_at, &spilled));
   ASSERT_TRUE(pinned_at.IsNil());
   rc->UpdateObjectPinnedAtRaylet(id, node_id);
-  ASSERT_TRUE(rc->IsPlasmaObjectPinned(id, &pinned_at));
+  ASSERT_TRUE(rc->IsPlasmaObjectPinnedOrSpilled(id, &pinned_at, &spilled));
   ASSERT_FALSE(pinned_at.IsNil());
 
   rc->RemoveLocalReference(id, nullptr);
-  ASSERT_FALSE(rc->IsPlasmaObjectPinned(id, &pinned_at));
+  ASSERT_FALSE(rc->IsPlasmaObjectPinnedOrSpilled(id, &pinned_at, &spilled));
   ASSERT_TRUE(deleted->count(id) > 0);
   deleted->clear();
 
@@ -2012,7 +2013,7 @@ TEST_F(ReferenceCountLineageEnabledTest, TestPlasmaLocation) {
   auto objects = rc->ResetObjectsOnRemovedNode(node_id);
   ASSERT_EQ(objects.size(), 1);
   ASSERT_EQ(objects[0], id);
-  ASSERT_TRUE(rc->IsPlasmaObjectPinned(id, &pinned_at));
+  ASSERT_TRUE(rc->IsPlasmaObjectPinnedOrSpilled(id, &pinned_at, &spilled));
   ASSERT_TRUE(pinned_at.IsNil());
   ASSERT_TRUE(deleted->count(id) > 0);
   deleted->clear();
@@ -2035,7 +2036,8 @@ TEST_F(ReferenceCountTest, TestFree) {
   ASSERT_EQ(deleted->count(id), 0);
   rc->UpdateObjectPinnedAtRaylet(id, node_id);
   NodeID pinned_at;
-  ASSERT_TRUE(rc->IsPlasmaObjectPinned(id, &pinned_at));
+  bool spilled;
+  ASSERT_TRUE(rc->IsPlasmaObjectPinnedOrSpilled(id, &pinned_at, &spilled));
   ASSERT_TRUE(pinned_at.IsNil());
   ASSERT_TRUE(rc->IsPlasmaObjectFreed(id));
   rc->RemoveLocalReference(id, nullptr);
@@ -2050,7 +2052,7 @@ TEST_F(ReferenceCountTest, TestFree) {
   rc->FreePlasmaObjects({id});
   ASSERT_TRUE(rc->IsPlasmaObjectFreed(id));
   ASSERT_TRUE(deleted->count(id) > 0);
-  ASSERT_TRUE(rc->IsPlasmaObjectPinned(id, &pinned_at));
+  ASSERT_TRUE(rc->IsPlasmaObjectPinnedOrSpilled(id, &pinned_at, &spilled));
   ASSERT_TRUE(pinned_at.IsNil());
   rc->RemoveLocalReference(id, nullptr);
   ASSERT_FALSE(rc->IsPlasmaObjectFreed(id));
diff --git a/src/ray/gcs/gcs_client/test/service_based_gcs_client_test.cc b/src/ray/gcs/gcs_client/test/service_based_gcs_client_test.cc
--- a/src/ray/gcs/gcs_client/test/service_based_gcs_client_test.cc
+++ b/src/ray/gcs/gcs_client/test/service_based_gcs_client_test.cc
@@ -435,7 +435,8 @@ class ServiceBasedGcsClientTest : public ::testing::Test {
 
   bool SubscribeToLocations(
       const ObjectID &object_id,
-      const gcs::SubscribeCallback<ObjectID, gcs::ObjectChangeNotification> &subscribe) {
+      const gcs::SubscribeCallback<ObjectID, std::vector<rpc::ObjectLocationChange>>
+          &subscribe) {
     std::promise<bool> promise;
     RAY_CHECK_OK(gcs_client_->Objects().AsyncSubscribeToLocations(
         object_id, subscribe,
@@ -474,9 +475,12 @@ class ServiceBasedGcsClientTest : public ::testing::Test {
     std::promise<bool> promise;
     std::vector<rpc::ObjectTableData> locations;
     RAY_CHECK_OK(gcs_client_->Objects().AsyncGetLocations(
-        object_id, [&locations, &promise](
-                       Status status, const std::vector<rpc::ObjectTableData> &result) {
-          locations = result;
+        object_id,
+        [&locations, &promise](Status status,
+                               const boost::optional<rpc::ObjectLocationInfo> &result) {
+          for (const auto &loc : result->locations()) {
+            locations.push_back(loc);
+          }
           promise.set_value(status.ok());
         }));
     EXPECT_TRUE(WaitReady(promise.get_future(), timeout_ms_));
@@ -852,11 +856,11 @@ TEST_F(ServiceBasedGcsClientTest, TestObjectInfo) {
   std::atomic<int> object_remove_count(0);
   auto on_subscribe = [&object_add_count, &object_remove_count](
                           const ObjectID &object_id,
-                          const gcs::ObjectChangeNotification &result) {
-    if (!result.GetData().empty()) {
-      if (result.IsAdded()) {
+                          const std::vector<rpc::ObjectLocationChange> &result) {
+    for (const auto &res : result) {
+      if (res.is_add()) {
         ++object_add_count;
-      } else if (result.IsRemoved()) {
+      } else {
         ++object_remove_count;
       }
     }
@@ -1020,16 +1024,18 @@ TEST_F(ServiceBasedGcsClientTest, TestObjectTableResubscribe) {
   std::atomic<int> object1_change_count(0);
   std::atomic<int> object2_change_count(0);
   ASSERT_TRUE(SubscribeToLocations(
-      object1_id, [&object1_change_count](const ObjectID &object_id,
-                                          const gcs::ObjectChangeNotification &result) {
-        if (!result.GetData().empty()) {
+      object1_id,
+      [&object1_change_count](const ObjectID &object_id,
+                              const std::vector<rpc::ObjectLocationChange> &result) {
+        if (!result.empty()) {
           ++object1_change_count;
         }
       }));
   ASSERT_TRUE(SubscribeToLocations(
-      object2_id, [&object2_change_count](const ObjectID &object_id,
-                                          const gcs::ObjectChangeNotification &result) {
-        if (!result.GetData().empty()) {
+      object2_id,
+      [&object2_change_count](const ObjectID &object_id,
+                              const std::vector<rpc::ObjectLocationChange> &result) {
+        if (!result.empty()) {
           ++object2_change_count;
         }
       }));
@@ -1241,8 +1247,8 @@ TEST_F(ServiceBasedGcsClientTest, TestMultiThreadSubAndUnsub) {
       for (int index = 0; index < sub_and_unsub_loop_count; ++index) {
         auto object_id = ObjectID::FromRandom();
         ASSERT_TRUE(SubscribeToLocations(
-            object_id,
-            [](const ObjectID &id, const gcs::ObjectChangeNotification &result) {}));
+            object_id, [](const ObjectID &id,
+                          const std::vector<rpc::ObjectLocationChange> &result) {}));
         gcs_client_->Objects().AsyncResubscribe(false);
         UnsubscribeToLocations(object_id);
       }
diff --git a/src/ray/gcs/gcs_server/test/gcs_server_rpc_test.cc b/src/ray/gcs/gcs_server/test/gcs_server_rpc_test.cc
--- a/src/ray/gcs/gcs_server/test/gcs_server_rpc_test.cc
+++ b/src/ray/gcs/gcs_server/test/gcs_server_rpc_test.cc
@@ -295,8 +295,8 @@ class GcsServerTest : public ::testing::Test {
         request, [&object_locations, &promise](
                      const Status &status, const rpc::GetObjectLocationsReply &reply) {
           RAY_CHECK_OK(status);
-          for (int index = 0; index < reply.object_table_data_list_size(); ++index) {
-            object_locations.push_back(reply.object_table_data_list(index));
+          for (const auto &loc : reply.location_info().locations()) {
+            object_locations.push_back(loc);
           }
           promise.set_value(true);
         });
diff --git a/src/ray/gcs/test/redis_object_info_accessor_test.cc b/src/ray/gcs/test/redis_object_info_accessor_test.cc
--- a/src/ray/gcs/test/redis_object_info_accessor_test.cc
+++ b/src/ray/gcs/test/redis_object_info_accessor_test.cc
@@ -70,9 +70,10 @@ TEST_F(RedisObjectInfoAccessorTest, TestGetAddRemove) {
     size_t total_size = elem.second.size();
     RAY_CHECK_OK(object_accessor.AsyncGetLocations(
         elem.first,
-        [this, total_size](Status status, const std::vector<ObjectTableData> &result) {
+        [this, total_size](Status status,
+                           const boost::optional<rpc::ObjectLocationInfo> &result) {
           RAY_CHECK_OK(status);
-          RAY_CHECK(total_size == result.size());
+          ASSERT_EQ(total_size, result->locations().size());
           --pending_count_;
         }));
   }
@@ -83,17 +84,18 @@ TEST_F(RedisObjectInfoAccessorTest, TestGetAddRemove) {
   // subscribe && delete
   // subscribe
   std::atomic<int> sub_pending_count(0);
-  auto subscribe = [this, &sub_pending_count](const ObjectID &object_id,
-                                              const ObjectChangeNotification &result) {
+  auto subscribe = [this, &sub_pending_count](
+                       const ObjectID &object_id,
+                       const std::vector<rpc::ObjectLocationChange> &result) {
     const auto it = object_id_to_data_.find(object_id);
     ASSERT_TRUE(it != object_id_to_data_.end());
     static size_t response_count = 1;
     size_t cur_count = response_count <= object_count_ ? copy_count_ : 1;
-    ASSERT_EQ(result.GetData().size(), cur_count);
-    rpc::GcsChangeMode change_mode = response_count <= object_count_
-                                         ? rpc::GcsChangeMode::APPEND_OR_ADD
-                                         : rpc::GcsChangeMode::REMOVE;
-    ASSERT_EQ(change_mode, result.GetGcsChangeMode());
+    ASSERT_EQ(result.size(), cur_count);
+    bool change_mode = response_count <= object_count_;
+    for (const auto &res : result) {
+      ASSERT_EQ(change_mode, res.is_add());
+    }
     ++response_count;
     --sub_pending_count;
   };
@@ -128,9 +130,10 @@ TEST_F(RedisObjectInfoAccessorTest, TestGetAddRemove) {
     size_t total_size = elem.second.size();
     RAY_CHECK_OK(object_accessor.AsyncGetLocations(
         elem.first,
-        [this, total_size](Status status, const std::vector<ObjectTableData> &result) {
+        [this, total_size](Status status,
+                           const boost::optional<rpc::ObjectLocationInfo> &result) {
           RAY_CHECK_OK(status);
-          ASSERT_EQ(total_size - 1, result.size());
+          ASSERT_EQ(total_size - 1, result->locations().size());
           --pending_count_;
         }));
   }
diff --git a/src/ray/object_manager/test/object_manager_stress_test.cc b/src/ray/object_manager/test/object_manager_stress_test.cc
--- a/src/ray/object_manager/test/object_manager_stress_test.cc
+++ b/src/ray/object_manager/test/object_manager_stress_test.cc
@@ -54,7 +54,8 @@ class MockServer {
         config_(object_manager_config),
         gcs_client_(gcs_client),
         object_manager_(main_service, node_id_, object_manager_config,
-                        std::make_shared<ObjectDirectory>(main_service, gcs_client_)) {
+                        std::make_shared<ObjectDirectory>(main_service, gcs_client_),
+                        nullptr) {
     RAY_CHECK_OK(RegisterGcs(main_service));
   }
 
diff --git a/src/ray/object_manager/test/object_manager_test.cc b/src/ray/object_manager/test/object_manager_test.cc
--- a/src/ray/object_manager/test/object_manager_test.cc
+++ b/src/ray/object_manager/test/object_manager_test.cc
@@ -50,7 +50,8 @@ class MockServer {
         config_(object_manager_config),
         gcs_client_(gcs_client),
         object_manager_(main_service, node_id_, object_manager_config,
-                        std::make_shared<ObjectDirectory>(main_service, gcs_client_)) {
+                        std::make_shared<ObjectDirectory>(main_service, gcs_client_),
+                        nullptr) {
     RAY_CHECK_OK(RegisterGcs(main_service));
   }
 
@@ -262,9 +263,9 @@ class TestObjectManager : public TestObjectManagerBase {
 
     RAY_CHECK_OK(server1->object_manager_.object_directory_->SubscribeObjectLocations(
         sub_id, object_1, rpc::Address(),
-        [this, sub_id, object_1, object_2](
-            const ray::ObjectID &object_id,
-            const std::unordered_set<ray::NodeID> &clients) {
+        [this, sub_id, object_1, object_2](const ray::ObjectID &object_id,
+                                           const std::unordered_set<ray::NodeID> &clients,
+                                           const std::string &spilled_url) {
           if (!clients.empty()) {
             TestWaitWhileSubscribed(sub_id, object_1, object_2);
           }
diff --git a/src/ray/raylet/reconstruction_policy_test.cc b/src/ray/raylet/reconstruction_policy_test.cc
--- a/src/ray/raylet/reconstruction_policy_test.cc
+++ b/src/ray/raylet/reconstruction_policy_test.cc
@@ -56,9 +56,9 @@ class MockObjectDirectory : public ObjectDirectoryInterface {
       const ObjectID object_id = callback.first;
       auto it = locations_.find(object_id);
       if (it == locations_.end()) {
-        callback.second(object_id, std::unordered_set<ray::NodeID>());
+        callback.second(object_id, std::unordered_set<ray::NodeID>(), "");
       } else {
-        callback.second(object_id, it->second);
+        callback.second(object_id, it->second, "");
       }
     }
     callbacks_.clear();