Call TriggerGlobalGC when the plasma store is full (#7337)

python/ray/_raylet.pyx

@@ -563,9 +563,10 @@ cdef void gc_collect() nogil:
         start = time.perf_counter()
         num_freed = gc.collect()
         end = time.perf_counter()
-        logger.info(
-            "gc.collect() freed {} refs in {} seconds".format(
-                num_freed, end - start))
+        if num_freed > 0:
+            logger.info(
+                "gc.collect() freed {} refs in {} seconds".format(
+                    num_freed, end - start))
 
 
 cdef shared_ptr[CBuffer] string_to_buffer(c_string& c_str):
@@ -1031,8 +1032,9 @@ cdef class CoreWorker:
                 contained_ids.push_back(
                     ObjectIDsToVector(serialized_object.contained_object_ids))
 
-        check_status(self.core_worker.get().AllocateReturnObjects(
-            return_ids, data_sizes, metadatas, contained_ids, returns))
+        with nogil:
+            check_status(self.core_worker.get().AllocateReturnObjects(
+                return_ids, data_sizes, metadatas, contained_ids, returns))
 
         for i, serialized_object in enumerate(serialized_objects):
             # A nullptr is returned if the object already exists.

python/ray/tests/test_reference_counting.py

@@ -116,6 +116,70 @@ def test_global_gc(shutdown_only):
         gc.enable()
 
 
+def test_global_gc_when_full(shutdown_only):
+    cluster = ray.cluster_utils.Cluster()
+    for _ in range(2):
+        cluster.add_node(
+            num_cpus=1, num_gpus=0, object_store_memory=100 * 1024 * 1024)
+    ray.init(address=cluster.address)
+
+    class LargeObjectWithCyclicRef:
+        def __init__(self):
+            self.loop = self
+            self.large_object = ray.put(
+                np.zeros(40 * 1024 * 1024, dtype=np.uint8))
+
+    @ray.remote(num_cpus=1)
+    class GarbageHolder:
+        def __init__(self):
+            gc.disable()
+            x = LargeObjectWithCyclicRef()
+            self.garbage = weakref.ref(x)
+
+        def has_garbage(self):
+            return self.garbage() is not None
+
+        def return_large_array(self):
+            return np.zeros(80 * 1024 * 1024, dtype=np.uint8)
+
+    try:
+        gc.disable()
+
+        # Local driver.
+        local_ref = weakref.ref(LargeObjectWithCyclicRef())
+
+        # Remote workers.
+        actors = [GarbageHolder.remote() for _ in range(2)]
+        assert local_ref() is not None
+        assert all(ray.get([a.has_garbage.remote() for a in actors]))
+
+        # GC should be triggered for all workers, including the local driver,
+        # when the driver tries to ray.put a value that doesn't fit in the
+        # object store. This should cause the captured ObjectIDs' numpy arrays
+        # to be evicted.
+        ray.put(np.zeros(80 * 1024 * 1024, dtype=np.uint8))
+        assert local_ref() is None
+        assert not any(ray.get([a.has_garbage.remote() for a in actors]))
+
+        # Local driver.
+        local_ref = weakref.ref(LargeObjectWithCyclicRef())
+
+        # Remote workers.
+        actors = [GarbageHolder.remote() for _ in range(2)]
+        assert local_ref() is not None
+        assert all(ray.get([a.has_garbage.remote() for a in actors]))
+
+        # GC should be triggered for all workers, including the local driver,
+        # when a remote task tries to put a return value that doesn't fit in
+        # the object store. This should cause the captured ObjectIDs' numpy
+        # arrays to be evicted.
+        ray.get(actors[0].return_large_array.remote())
+        assert local_ref() is None
+        assert not any(ray.get([a.has_garbage.remote() for a in actors]))
+    finally:
+        gc.enable()
+
+
 def test_local_refcounts(ray_start_regular):
     oid1 = ray.put(None)
     check_refcounts({oid1: (1, 0)})

src/ray/core_worker/core_worker.cc

@@ -181,7 +181,8 @@ CoreWorker::CoreWorker(const WorkerType worker_type, const Language language,
   io_thread_ = std::thread(&CoreWorker::RunIOService, this);
 
   plasma_store_provider_.reset(new CoreWorkerPlasmaStoreProvider(
-      store_socket, local_raylet_client_, check_signals_));
+      store_socket, local_raylet_client_, check_signals_,
+      boost::bind(&CoreWorker::TriggerGlobalGC, this)));
   memory_store_.reset(new CoreWorkerMemoryStore(
       [this](const RayObject &obj, const ObjectID &obj_id) {
         RAY_LOG(DEBUG) << "Promoting object to plasma " << obj_id;

src/ray/core_worker/store_provider/plasma_store_provider.cc

@@ -10,9 +10,10 @@ namespace ray {
 CoreWorkerPlasmaStoreProvider::CoreWorkerPlasmaStoreProvider(
     const std::string &store_socket,
     const std::shared_ptr<raylet::RayletClient> raylet_client,
-    std::function<Status()> check_signals)
+    std::function<Status()> check_signals, std::function<void()> on_store_full)
     : raylet_client_(raylet_client) {
   check_signals_ = check_signals;
+  on_store_full_ = on_store_full;
   RAY_ARROW_CHECK_OK(store_client_.Connect(store_socket));
 }
 
@@ -76,13 +77,19 @@ Status CoreWorkerPlasmaStoreProvider::Create(const std::shared_ptr<Buffer> &meta
               << "is full. Object size is " << data_size << " bytes.";
       status = Status::ObjectStoreFull(message.str());
       if (max_retries < 0 || retries < max_retries) {
-        RAY_LOG(ERROR) << message.str() << " Plasma store status:\n"
-                       << MemoryUsageString() << "\nWaiting " << delay
+        RAY_LOG(ERROR) << message.str() << "\nWaiting " << delay
                        << "ms for space to free up...";
+        if (on_store_full_) {
+          on_store_full_();
+        }
         usleep(1000 * delay);
         delay *= 2;
         retries += 1;
         should_retry = true;
+      } else {
+        RAY_LOG(ERROR) << "Failed to put object " << object_id << " after " << max_retries
+                       << " attempts. Plasma store status:\n"
+                       << MemoryUsageString();
       }
     } else if (plasma::IsPlasmaObjectExists(plasma_status)) {
       RAY_LOG(WARNING) << "Trying to put an object that already existed in plasma: "
@@ -205,9 +212,9 @@ Status CoreWorkerPlasmaStoreProvider::Get(
     return Status::OK();
   }
 
-  // If not all objects were successfully fetched, repeatedly call FetchOrReconstruct and
-  // Get from the local object store in batches. This loop will run indefinitely until the
-  // objects are all fetched if timeout is -1.
+  // If not all objects were successfully fetched, repeatedly call FetchOrReconstruct
+  // and Get from the local object store in batches. This loop will run indefinitely
+  // until the objects are all fetched if timeout is -1.
   int unsuccessful_attempts = 0;
   bool should_break = false;
   bool timed_out = false;
@@ -346,7 +353,8 @@ void CoreWorkerPlasmaStoreProvider::WarnIfAttemptedTooManyTimes(
     RAY_LOG(WARNING)
         << "Attempted " << num_attempts << " times to reconstruct objects, but "
         << "some objects are still unavailable. If this message continues to print,"
-        << " it may indicate that object's creating task is hanging, or something wrong"
+        << " it may indicate that object's creating task is hanging, or something "
+           "wrong"
         << " happened in raylet backend. " << remaining.size()
         << " object(s) pending: " << oss.str() << ".";
   }

src/ray/core_worker/store_provider/plasma_store_provider.h

@@ -21,7 +21,8 @@ class CoreWorkerPlasmaStoreProvider {
  public:
   CoreWorkerPlasmaStoreProvider(const std::string &store_socket,
                                 const std::shared_ptr<raylet::RayletClient> raylet_client,
-                                std::function<Status()> check_signals);
+                                std::function<Status()> check_signals,
+                                std::function<void()> on_store_full = nullptr);
 
   ~CoreWorkerPlasmaStoreProvider();
 
@@ -119,6 +120,7 @@ class CoreWorkerPlasmaStoreProvider {
   plasma::PlasmaClient store_client_;
   std::mutex store_client_mutex_;
   std::function<Status()> check_signals_;
+  std::function<void()> on_store_full_;
 };
 
 }  // namespace ray