Add internal_api.global_gc() method, which triggers gc.collect() on all workers (#7327)

python/ray/_raylet.pyx

@@ -562,6 +562,16 @@ cdef CRayStatus check_signals() nogil:
     return CRayStatus.OK()
 
 
+cdef void gc_collect() nogil:
+    with gil:
+        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))
+
+
 cdef shared_ptr[CBuffer] string_to_buffer(c_string& c_str):
     cdef shared_ptr[CBuffer] empty_metadata
     if c_str.size() == 0:
@@ -607,7 +617,7 @@ cdef class CoreWorker:
             raylet_socket.encode("ascii"), job_id.native(),
             gcs_options.native()[0], log_dir.encode("utf-8"),
             node_ip_address.encode("utf-8"), node_manager_port,
-            task_execution_handler, check_signals, True))
+            task_execution_handler, check_signals, gc_collect, True))
 
     def run_task_loop(self):
         with nogil:
@@ -760,6 +770,10 @@ cdef class CoreWorker:
             check_status(self.core_worker.get().Delete(
                 free_ids, local_only, delete_creating_tasks))
 
+    def global_gc(self):
+        with nogil:
+            self.core_worker.get().TriggerGlobalGC()
+
     def set_object_store_client_options(self, client_name,
                                         int64_t limit_bytes):
         try:

python/ray/includes/libcoreworker.pxd

@@ -97,6 +97,7 @@ cdef extern from "ray/core_worker/core_worker.h" nogil:
                         c_vector[shared_ptr[CRayObject]] *returns,
                         const CWorkerID &worker_id) nogil,
                     CRayStatus() nogil,
+                    void() nogil,
                     c_bool ref_counting_enabled)
         CWorkerType &GetWorkerType()
         CLanguage &GetLanguage()
@@ -173,6 +174,7 @@ cdef extern from "ray/core_worker/core_worker.h" nogil:
                         int64_t timeout_ms, c_vector[c_bool] *results)
         CRayStatus Delete(const c_vector[CObjectID] &object_ids,
                           c_bool local_only, c_bool delete_creating_tasks)
+        CRayStatus TriggerGlobalGC()
         c_string MemoryUsageString()
 
         CWorkerContext &GetWorkerContext()

python/ray/internal/internal_api.py

@@ -1,7 +1,14 @@
 import ray.worker
 from ray import profiling
 
-__all__ = ["free"]
+__all__ = ["free", "global_gc"]
+
+
+def global_gc():
+    """Trigger gc.collect() on all workers in the cluster."""
+
+    worker = ray.worker.get_global_worker()
+    worker.core_worker.global_gc()
 
 
 def free(object_ids, local_only=False, delete_creating_tasks=False):

python/ray/tests/test_reference_counting.py

@@ -3,9 +3,11 @@ import copy
 import json
 import logging
 import os
+import gc
 import tempfile
 import time
 import uuid
+import weakref
 
 import numpy as np
 
@@ -14,6 +16,7 @@ import pytest
 import ray
 import ray.cluster_utils
 import ray.test_utils
+from ray.internal.internal_api import global_gc
 
 logger = logging.getLogger(__name__)
 
@@ -72,6 +75,46 @@ def check_refcounts(expected, timeout=10):
                 time.sleep(0.1)
 
 
+def test_global_gc(shutdown_only):
+    cluster = ray.cluster_utils.Cluster()
+    for _ in range(2):
+        cluster.add_node(num_cpus=1, num_gpus=0)
+    ray.init(address=cluster.address)
+
+    class ObjectWithCyclicRef:
+        def __init__(self):
+            self.loop = self
+
+    @ray.remote(num_cpus=1)
+    class GarbageHolder:
+        def __init__(self):
+            gc.disable()
+            x = ObjectWithCyclicRef()
+            self.garbage = weakref.ref(x)
+
+        def has_garbage(self):
+            return self.garbage() is not None
+
+    try:
+        gc.disable()
+
+        # Local driver.
+        local_ref = weakref.ref(ObjectWithCyclicRef())
+
+        # 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.
+        global_gc()
+        time.sleep(1)
+        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)})