gc collect from a gc_thread (ray-project#55838)

Ray proactively triggers gc.collect() on idle workers to release Python
objects that may still hold Plasma shared memory (shm) references.
In the current implementation in (_raylet.pyx gc_collect()), Ray calls
gc.collect() from Cython under a with gil block periodically.
If the Python object graph is complex (e.g., cyclic references with
finalizers), gc.collect() may take a long time. During this period,
since the GIL is held for the entire collection, user code is completely
frozen if gc.collect() time is longer than the periodic interval (e.g.,
10s).

We propose decoupling GC execution from the RPC call:

gc_collect in Cython should not directly run gc.collect().
Instead, it should "signal an event" with minimum execution time (e.g.,
using a threading.Event or similar).
A dedicated Python GC thread consumes this event and executes
gc.collect() asynchronously, with a configurable GC interval.

## Related issue number

Closes ray-project#55837

---------

Signed-off-by: Mao Yancan <[email protected]>
Co-authored-by: Mao Yancan <[email protected]>
Co-authored-by: Edward Oakes <[email protected]>
Signed-off-by: zac <[email protected]>

Author: 3 people

python/ray/_private/gc_collect_manager.py

@@ -0,0 +1,65 @@
+import gc
+import logging
+import threading
+import time
+from typing import Callable, Optional
+
+logger = logging.getLogger(__name__)
+
+
+class PythonGCThread(threading.Thread):
+    """A background thread that triggers Python garbage collection.
+
+    This thread waits for GC events from CoreWorker and triggers `gc.collect()` when
+    requested, ensuring that collections are spaced out by at least
+    `min_interval_s` seconds."""
+
+    def __init__(
+        self, *, min_interval_s: int = 5, gc_collect_func: Optional[Callable] = None
+    ):
+        logger.debug("Starting Python GC thread")
+        super().__init__(name="PythonGCThread", daemon=True)
+        self._should_exit = False
+        self._last_gc_time = float("-inf")
+        self._min_gc_interval = min_interval_s
+        self._gc_event = threading.Event()
+        # Set the gc_collect_func for UT, defaulting to gc.collect if None
+        self._gc_collect_func = gc_collect_func or gc.collect
+
+    def trigger_gc(self) -> None:
+        self._gc_event.set()
+
+    def run(self):
+        while not self._should_exit:
+            self._gc_event.wait()
+            self._gc_event.clear()
+
+            if self._should_exit:
+                break
+
+            time_since_last_gc = time.monotonic() - self._last_gc_time
+            if time_since_last_gc < self._min_gc_interval:
+                logger.debug(
+                    f"Skipping GC, only {time_since_last_gc:.2f}s since last GC"
+                )
+                continue
+
+            try:
+                start = time.monotonic()
+                num_freed = self._gc_collect_func()
+                self._last_gc_time = time.monotonic()
+                if num_freed > 0:
+                    logger.debug(
+                        "gc.collect() freed {} refs in {} seconds".format(
+                            num_freed, self._last_gc_time - start
+                        )
+                    )
+            except Exception as e:
+                logger.error(f"Error during GC: {e}")
+                self._last_gc_time = time.monotonic()
+
+    def stop(self):
+        logger.debug("Stopping Python GC thread")
+        self._should_exit = True
+        self._gc_event.set()
+        self.join()

python/ray/_private/ray_constants.py

@@ -584,3 +584,5 @@ def gcs_actor_scheduling_enabled():
 FETCH_FAIL_TIMEOUT_SECONDS = (
     env_integer("RAY_fetch_fail_timeout_milliseconds", 60000) / 1000
 )
+
+RAY_GC_MIN_COLLECT_INTERVAL = env_float("RAY_GC_MIN_COLLECT_INTERVAL_S", 5)

python/ray/_raylet.pxd

@@ -140,6 +140,7 @@ cdef class CoreWorker:
         object _task_id_to_future_lock
         dict _task_id_to_future
         object event_loop_executor
+        object _gc_thread
 
     cdef unique_ptr[CAddress] _convert_python_address(self, address=*)
     cdef store_task_output(

python/ray/_raylet.pyx

@@ -248,6 +248,7 @@ from ray._private.utils import DeferSigint
 from ray._private.object_ref_generator import DynamicObjectRefGenerator
 from ray.util.annotations import PublicAPI
 from ray._private.custom_types import TensorTransportEnum
+from ray._private.gc_collect_manager import PythonGCThread
 
 # Expose GCC & Clang macro to report
 # whether C++ optimizations were enabled during compilation.
@@ -2496,14 +2497,21 @@ cdef CRayStatus check_signals() nogil:
 
 
 cdef void gc_collect(c_bool triggered_by_global_gc) nogil:
-    with gil:
-        start = time.perf_counter()
-        num_freed = gc.collect()
-        end = time.perf_counter()
-        if num_freed > 0:
-            logger.debug(
-                "gc.collect() freed {} refs in {} seconds".format(
-                    num_freed, end - start))
+     with gil:
+        if RayConfig.instance().start_python_gc_manager_thread():
+            start = time.perf_counter()
+            worker = ray._private.worker.global_worker
+            worker.core_worker.trigger_gc()
+            end = time.perf_counter()
+            logger.debug("GC event triggered in {} seconds".format(end - start))
+        else:
+            start = time.perf_counter()
+            num_freed = gc.collect()
+            end = time.perf_counter()
+            if num_freed > 0:
+                logger.debug(
+                    "gc.collect() freed {} refs in {} seconds".format(
+                        num_freed, end - start))
 
 
 cdef c_vector[c_string] spill_objects_handler(
@@ -3054,13 +3062,21 @@ cdef class CoreWorker:
         self._task_id_to_future = {}
         self.event_loop_executor = None
 
+        self._gc_thread = None
+        if RayConfig.instance().start_python_gc_manager_thread():
+            self._gc_thread = PythonGCThread(min_interval_s=ray_constants.RAY_GC_MIN_COLLECT_INTERVAL)
+            self._gc_thread.start()
+
     def shutdown_driver(self):
         # If it's a worker, the core worker process should have been
         # shutdown. So we can't call
         # `CCoreWorkerProcess.GetCoreWorker().GetWorkerType()` here.
         # Instead, we use the cached `is_driver` flag to test if it's a
         # driver.
         assert self.is_driver
+        if self._gc_thread is not None:
+            self._gc_thread.stop()
+            self._gc_thread = None
         with nogil:
             CCoreWorkerProcess.Shutdown()
 
@@ -4719,6 +4735,9 @@ cdef class CoreWorker:
 
         return self.current_runtime_env
 
+    def trigger_gc(self):
+        self._gc_thread.trigger_gc()
+
     def get_pending_children_task_ids(self, parent_task_id: TaskID):
         cdef:
             CTaskID c_parent_task_id = parent_task_id.native()

python/ray/includes/ray_config.pxd

@@ -86,3 +86,5 @@ cdef extern from "ray/common/ray_config.h" nogil:
         int maximum_gcs_destroyed_actor_cached_count() const
 
         c_bool record_task_actor_creation_sites() const
+
+        c_bool start_python_gc_manager_thread() const

python/ray/includes/ray_config.pxi

@@ -140,3 +140,7 @@ cdef class Config:
     @staticmethod
     def maximum_gcs_destroyed_actor_cached_count():
         return RayConfig.instance().maximum_gcs_destroyed_actor_cached_count()
+
+    @staticmethod
+    def start_python_gc_manager_thread():
+        return RayConfig.instance().start_python_gc_manager_thread()

python/ray/tests/test_basic.py

@@ -267,7 +267,7 @@ def get_thread_count(self):
         ray.get(actor.get_thread_count.remote())
     # Lowering these numbers in this assert should be celebrated,
     # increasing these numbers should be scrutinized
-    assert ray.get(actor.get_thread_count.remote()) in {24, 25}
+    assert ray.get(actor.get_thread_count.remote()) in {24, 25, 26}
 
 
 # https://github.com/ray-project/ray/issues/7287

python/ray/tests/test_global_gc.py

@@ -2,14 +2,17 @@
 import gc
 import logging
 import sys
+import time
 import weakref
+from unittest.mock import Mock
 
 import numpy as np
 import pytest
 
 import ray
 import ray.cluster_utils
 from ray._common.test_utils import wait_for_condition
+from ray._private.gc_collect_manager import PythonGCThread
 from ray._private.internal_api import global_gc
 
 logger = logging.getLogger(__name__)
@@ -216,5 +219,136 @@ def f(self):
         gc.enable()
 
 
+def test_local_gc_called_once_per_interval(shutdown_only):
+    ray.init(
+        num_cpus=2,
+        _system_config={
+            "local_gc_interval_s": 1,
+            "local_gc_min_interval_s": 0,
+            "global_gc_min_interval_s": 0,
+        },
+    )
+
+    class ObjectWithCyclicRef:
+        def __init__(self):
+            self.loop = self
+
+    @ray.remote(num_cpus=1)
+    class GarbageHolder:
+        def __init__(self):
+            gc.disable()
+            self.garbage = None
+
+        def make_garbage(self):
+            x = ObjectWithCyclicRef()
+            self.garbage = weakref.ref(x)
+            return True
+
+        def has_garbage(self):
+            return self.garbage() is not None
+
+    def all_garbage_collected(local_ref):
+        return local_ref() is None and not any(
+            ray.get([a.has_garbage.remote() for a in actors])
+        )
+
+    try:
+        gc.disable()
+
+        # Round 1: first batch of garbage should be collected
+        # Local driver.
+        local_ref = weakref.ref(ObjectWithCyclicRef())
+        # Remote workers.
+        actors = [GarbageHolder.remote() for _ in range(2)]
+        ray.get([a.make_garbage.remote() for a in actors])
+
+        assert local_ref() is not None
+        assert all(ray.get([a.has_garbage.remote() for a in actors]))
+
+        wait_for_condition(
+            lambda: all_garbage_collected(local_ref),
+        )
+
+        # Round 2: second batch should NOT be collected within min_interval
+        local_ref = weakref.ref(ObjectWithCyclicRef())
+        ray.get([a.make_garbage.remote() for a in actors])
+
+        with pytest.raises(RuntimeError):
+            wait_for_condition(
+                lambda: all_garbage_collected(local_ref),
+                timeout=2.0,  # shorter than min_interval
+                retry_interval_ms=50,
+            )
+
+        # Round 3: after min_interval passes, garbage should be collected
+        wait_for_condition(
+            lambda: all_garbage_collected(local_ref),
+            timeout=10.0,
+            retry_interval_ms=50,
+        )
+
+    finally:
+        gc.enable()
+
+
+def test_gc_manager_thread_basic_functionality():
+    mock_gc_collect = Mock(return_value=10)
+
+    gc_thread = PythonGCThread(min_interval_s=1, gc_collect_func=mock_gc_collect)
+
+    try:
+        gc_thread.start()
+        assert gc_thread.is_alive()
+
+        gc_thread.trigger_gc()
+
+        wait_for_condition(lambda: mock_gc_collect.call_count == 1, timeout=2)
+
+        mock_gc_collect.assert_called_once()
+
+    finally:
+        gc_thread.stop()
+        assert not gc_thread.is_alive()
+
+
+def test_gc_manager_thread_min_interval_throttling():
+    mock_gc_collect = Mock(return_value=5)
+
+    gc_thread = PythonGCThread(min_interval_s=2, gc_collect_func=mock_gc_collect)
+
+    try:
+        gc_thread.start()
+
+        for _ in range(3):
+            gc_thread.trigger_gc()
+            time.sleep(1)
+
+        wait_for_condition(lambda: mock_gc_collect.call_count == 2, timeout=2)
+
+        assert mock_gc_collect.call_count == 2
+
+    finally:
+        gc_thread.stop()
+
+
+def test_gc_manager_thread_exception_handling():
+    mock_gc_collect = Mock(side_effect=RuntimeError("GC failed"))
+
+    gc_thread = PythonGCThread(min_interval_s=5, gc_collect_func=mock_gc_collect)
+
+    try:
+        gc_thread.start()
+
+        for _ in range(3):
+            gc_thread.trigger_gc()
+            time.sleep(0.1)
+
+        assert gc_thread.is_alive()
+        mock_gc_collect.assert_called_once()
+
+    finally:
+        gc_thread.stop()
+
+
 if __name__ == "__main__":
     sys.exit(pytest.main(["-sv", __file__]))

src/ray/common/ray_config_def.h

@@ -945,3 +945,6 @@ RAY_CONFIG(int32_t, raylet_rpc_server_reconnect_timeout_s, 60)
 // process getting spawned.  Setting to zero or less maintains the default
 // number of threads grpc will spawn.
 RAY_CONFIG(int64_t, worker_num_grpc_internal_threads, 0)
+
+// Whether to start a background thread to manage Python GC in workers.
+RAY_CONFIG(bool, start_python_gc_manager_thread, true)

src/ray/raylet/worker_pool.cc

@@ -461,6 +461,7 @@ WorkerPool::BuildProcessCommandArgs(const Language &language,
     // Support forking in gRPC.
     env.insert({"GRPC_ENABLE_FORK_SUPPORT", "True"});
     env.insert({"GRPC_POLL_STRATEGY", "poll"});
+    env.insert({"RAY_start_python_gc_manager_thread", "0"});
   }
 
   return {std::move(worker_command_args), std::move(env)};