Remove ref counting dependencies on ray.get() (#6412)

* Remove ref counting dependencies on Get()

* comment

* don't send IDs when disabled

* pass through internal config

* fix

* allow reinit

* remove flag

python/ray/_raylet.pyx

@@ -138,6 +138,20 @@ else:
 if PY3:
     from ray.async_compat import sync_to_async
 
+
+def set_internal_config(dict options):
+    cdef:
+        unordered_map[c_string, c_string] c_options
+
+    if options is None:
+        return
+
+    for key, value in options.items():
+        c_options[str(key).encode("ascii")] = str(value).encode("ascii")
+
+    RayConfig.instance().initialize(c_options)
+
+
 cdef int check_status(const CRayStatus& status) nogil except -1:
     if status.ok():
         return 0

python/ray/includes/ray_config.pxd

@@ -86,4 +86,4 @@ cdef extern from "ray/common/ray_config.h" nogil:
 
         uint32_t maximum_gcs_deletion_batch_size() const
 
-        void initialize(const unordered_map[c_string, int] &config_map)
+        void initialize(const unordered_map[c_string, c_string] &config_map)

python/ray/services.py

@@ -1199,10 +1199,12 @@ def start_raylet(redis_address,
                             "--object-store-name={} "
                             "--raylet-name={} "
                             "--redis-address={} "
+                            "--config-list={} "
                             "--temp-dir={}".format(
                                 sys.executable, worker_path, node_ip_address,
                                 node_manager_port, plasma_store_name,
-                                raylet_name, redis_address, temp_dir))
+                                raylet_name, redis_address, config_str,
+                                temp_dir))
     if redis_password:
         start_worker_command += " --redis-password {}".format(redis_password)
 

python/ray/worker.py

@@ -795,7 +795,9 @@ def init(address=None,
         log_to_driver=log_to_driver,
         worker=global_worker,
         driver_object_store_memory=driver_object_store_memory,
-        job_id=job_id)
+        job_id=job_id,
+        internal_config=json.loads(_internal_config)
+        if _internal_config else {})
 
     for hook in _post_init_hooks:
         hook()
@@ -1064,7 +1066,8 @@ def connect(node,
             log_to_driver=False,
             worker=global_worker,
             driver_object_store_memory=None,
-            job_id=None):
+            job_id=None,
+            internal_config=None):
     """Connect this worker to the raylet, to Plasma, and to Redis.
 
     Args:
@@ -1077,6 +1080,8 @@ def connect(node,
         driver_object_store_memory: Limit the amount of memory the driver can
             use in the object store when creating objects.
         job_id: The ID of job. If it's None, then we will generate one.
+        internal_config: Dictionary of (str,str) containing internal config
+            options to override the defaults.
     """
     # Do some basic checking to make sure we didn't call ray.init twice.
     error_message = "Perhaps you called ray.init twice by accident?"
@@ -1087,6 +1092,8 @@ def connect(node,
     if not faulthandler.is_enabled():
         faulthandler.enable(all_threads=False)
 
+    ray._raylet.set_internal_config(internal_config)
+
     if mode is not LOCAL_MODE:
         # Create a Redis client to primary.
         # The Redis client can safely be shared between threads. However,

python/ray/workers/default_worker.py

@@ -3,6 +3,7 @@ from __future__ import division
 from __future__ import print_function
 
 import argparse
+import json
 
 import ray
 import ray.actor
@@ -55,6 +56,12 @@ parser.add_argument(
     type=str,
     default=ray_constants.LOGGER_FORMAT,
     help=ray_constants.LOGGER_FORMAT_HELP)
+parser.add_argument(
+    "--config-list",
+    required=False,
+    type=str,
+    default=None,
+    help="Override internal config options for the worker process.")
 parser.add_argument(
     "--temp-dir",
     required=False,
@@ -77,6 +84,15 @@ if __name__ == "__main__":
 
     ray.utils.setup_logger(args.logging_level, args.logging_format)
 
+    internal_config = {}
+    if args.config_list is not None:
+        config_list = args.config_list.split(",")
+        if len(config_list) > 1:
+            i = 0
+            while i < len(config_list):
+                internal_config[config_list[i]] = config_list[i + 1]
+                i += 2
+
     ray_params = RayParams(
         node_ip_address=args.node_ip_address,
         node_manager_port=args.node_manager_port,
@@ -86,7 +102,9 @@ if __name__ == "__main__":
         raylet_socket_name=args.raylet_name,
         temp_dir=args.temp_dir,
         load_code_from_local=args.load_code_from_local,
-        use_pickle=args.use_pickle)
+        use_pickle=args.use_pickle,
+        _internal_config=json.dumps(internal_config),
+    )
 
     node = ray.node.Node(
         ray_params,
@@ -95,5 +113,6 @@ if __name__ == "__main__":
         spawn_reaper=False,
         connect_only=True)
     ray.worker._global_node = node
-    ray.worker.connect(node, mode=ray.WORKER_MODE)
+    ray.worker.connect(
+        node, mode=ray.WORKER_MODE, internal_config=internal_config)
     ray.worker.global_worker.main_loop()

src/ray/common/ray_config.h

@@ -43,21 +43,15 @@ class RayConfig {
   }
 
   void initialize(const std::unordered_map<std::string, std::string> &config_map) {
-    RAY_CHECK(!initialized_);
     for (auto const &pair : config_map) {
       // We use a big chain of if else statements because C++ doesn't allow
       // switch statements on strings.
 #include "ray_config_def.h"
       RAY_LOG(FATAL) << "Received unexpected config parameter " << pair.first;
     }
-    initialized_ = true;
   }
 /// ---------------------------------------------------------------------
 #undef RAY_CONFIG
-
-  /// Whether the initialization of the instance has been called before.
-  /// The RayConfig instance can only (and must) be initialized once.
-  bool initialized_ = false;
 };
 // clang-format on
 

src/ray/core_worker/core_worker.cc

@@ -1,11 +1,11 @@
+#include "ray/core_worker/core_worker.h"
+
 #include <cstdlib>
 
 #include "boost/fiber/all.hpp"
-
 #include "ray/common/ray_config.h"
 #include "ray/common/task/task_util.h"
 #include "ray/core_worker/context.h"
-#include "ray/core_worker/core_worker.h"
 #include "ray/core_worker/transport/direct_actor_transport.h"
 #include "ray/core_worker/transport/raylet_transport.h"
 
@@ -261,14 +261,16 @@ void CoreWorker::SetCurrentTaskId(const TaskID &task_id) {
 }
 
 void CoreWorker::ReportActiveObjectIDs() {
-  std::unordered_set<ObjectID> active_object_ids =
-      reference_counter_->GetAllInScopeObjectIDs();
-  RAY_LOG(DEBUG) << "Sending " << active_object_ids.size() << " object IDs to raylet.";
-  auto max_active = RayConfig::instance().raylet_max_active_object_ids();
-  if (max_active && active_object_ids.size() > max_active) {
-    RAY_LOG(INFO) << active_object_ids.size() << " object IDs are currently in scope.";
+  std::unordered_set<ObjectID> active_object_ids;
+  size_t max_active = RayConfig::instance().raylet_max_active_object_ids();
+  if (max_active > 0) {
+    active_object_ids = reference_counter_->GetAllInScopeObjectIDs();
+    if (active_object_ids.size() > max_active) {
+      RAY_LOG(INFO) << active_object_ids.size() << " object IDs are currently in scope.";
+    }
   }
 
+  RAY_LOG(DEBUG) << "Sending " << active_object_ids.size() << " object IDs to raylet.";
   if (!raylet_client_->ReportActiveObjectIDs(active_object_ids).ok()) {
     RAY_LOG(ERROR) << "Raylet connection failed. Shutting down.";
     Shutdown();
@@ -402,6 +404,12 @@ Status CoreWorker::Get(const std::vector<ObjectID> &ids, const int64_t timeout_m
         // object.
         will_throw_exception = true;
       }
+      // If we got the result for this plasma ObjectID, the task that created it must
+      // have finished. Therefore, we can safely remove its reference counting
+      // dependencies.
+      if (!ids[i].IsDirectCallType()) {
+        RemoveObjectIDDependencies(ids[i]);
+      }
     } else {
       missing_result = true;
     }

src/ray/core_worker/core_worker.h

@@ -2,7 +2,6 @@
 #define RAY_CORE_WORKER_CORE_WORKER_H
 
 #include "absl/container/flat_hash_map.h"
-
 #include "ray/common/buffer.h"
 #include "ray/core_worker/actor_handle.h"
 #include "ray/core_worker/common.h"
@@ -435,6 +434,17 @@ class CoreWorker {
                               std::vector<std::shared_ptr<RayObject>> *args,
                               std::vector<ObjectID> *arg_reference_ids);
 
+  /// Remove reference counting dependencies of this object ID.
+  ///
+  /// \param[in] object_id The object whose dependencies should be removed.
+  void RemoveObjectIDDependencies(const ObjectID &object_id) {
+    std::vector<ObjectID> deleted;
+    reference_counter_->RemoveDependencies(object_id, &deleted);
+    if (ref_counting_enabled_) {
+      memory_store_->Delete(deleted);
+    }
+  }
+
   /// Type of this worker (i.e., DRIVER or WORKER).
   const WorkerType worker_type_;
 

src/ray/core_worker/reference_count.cc

@@ -42,6 +42,23 @@ void ReferenceCounter::RemoveReference(const ObjectID &object_id,
   RemoveReferenceRecursive(object_id, deleted);
 }
 
+void ReferenceCounter::RemoveDependencies(const ObjectID &object_id,
+                                          std::vector<ObjectID> *deleted) {
+  absl::MutexLock lock(&mutex_);
+  auto entry = object_id_refs_.find(object_id);
+  if (entry == object_id_refs_.end()) {
+    RAY_LOG(WARNING) << "Tried to remove dependencies for nonexistent object ID: "
+                     << object_id;
+    return;
+  }
+  if (entry->second.second) {
+    for (const ObjectID &pending_task_object_id : *entry->second.second) {
+      RemoveReferenceRecursive(pending_task_object_id, deleted);
+    }
+    entry->second.second = nullptr;
+  }
+}
+
 void ReferenceCounter::RemoveReferenceRecursive(const ObjectID &object_id,
                                                 std::vector<ObjectID> *deleted) {
   auto entry = object_id_refs_.find(object_id);

src/ray/core_worker/reference_count.h

@@ -4,7 +4,6 @@
 #include "absl/base/thread_annotations.h"
 #include "absl/container/flat_hash_map.h"
 #include "absl/synchronization/mutex.h"
-
 #include "ray/common/id.h"
 #include "ray/util/logging.h"
 
@@ -38,6 +37,14 @@ class ReferenceCounter {
                        std::shared_ptr<std::vector<ObjectID>> dependencies)
       LOCKS_EXCLUDED(mutex_);
 
+  /// Remove any references to dependencies that this object may have. This does *not*
+  /// decrease the object's own reference count.
+  ///
+  /// \param[in] object_id The object whose dependencies should be removed.
+  /// \param[out] deleted List to store objects that hit zero ref count.
+  void RemoveDependencies(const ObjectID &object_id, std::vector<ObjectID> *deleted)
+      LOCKS_EXCLUDED(mutex_);
+
   /// Returns the total number of ObjectIDs currently in scope.
   size_t NumObjectIDsInScope() const LOCKS_EXCLUDED(mutex_);
 

src/ray/core_worker/reference_count_test.cc

@@ -1,8 +1,9 @@
+#include "ray/core_worker/reference_count.h"
+
 #include <vector>
 
 #include "gtest/gtest.h"
 #include "ray/common/ray_object.h"
-#include "ray/core_worker/reference_count.h"
 #include "ray/core_worker/store_provider/memory_store/memory_store.h"
 
 namespace ray {
@@ -141,6 +142,49 @@ TEST_F(ReferenceCountTest, TestRecursiveDependencies) {
   ASSERT_EQ(out.size(), 4);
 }
 
+TEST_F(ReferenceCountTest, TestRemoveDependenciesOnly) {
+  std::vector<ObjectID> out;
+  ObjectID id1 = ObjectID::FromRandom();
+  ObjectID id2 = ObjectID::FromRandom();
+  ObjectID id3 = ObjectID::FromRandom();
+  ObjectID id4 = ObjectID::FromRandom();
+
+  std::shared_ptr<std::vector<ObjectID>> deps2 =
+      std::make_shared<std::vector<ObjectID>>();
+  deps2->push_back(id3);
+  deps2->push_back(id4);
+  rc->AddOwnedObject(id2, TaskID::Nil(), rpc::Address(), deps2);
+
+  std::shared_ptr<std::vector<ObjectID>> deps1 =
+      std::make_shared<std::vector<ObjectID>>();
+  deps1->push_back(id2);
+  rc->AddOwnedObject(id1, TaskID::Nil(), rpc::Address(), deps1);
+
+  rc->AddLocalReference(id1);
+  rc->AddLocalReference(id2);
+  rc->AddLocalReference(id4);
+  ASSERT_EQ(rc->NumObjectIDsInScope(), 4);
+
+  rc->RemoveDependencies(id2, &out);
+  ASSERT_EQ(rc->NumObjectIDsInScope(), 3);
+  ASSERT_EQ(out.size(), 1);
+  rc->RemoveDependencies(id1, &out);
+  ASSERT_EQ(rc->NumObjectIDsInScope(), 3);
+  ASSERT_EQ(out.size(), 1);
+
+  rc->RemoveLocalReference(id1, &out);
+  ASSERT_EQ(rc->NumObjectIDsInScope(), 2);
+  ASSERT_EQ(out.size(), 2);
+
+  rc->RemoveLocalReference(id2, &out);
+  ASSERT_EQ(rc->NumObjectIDsInScope(), 1);
+  ASSERT_EQ(out.size(), 3);
+
+  rc->RemoveLocalReference(id4, &out);
+  ASSERT_EQ(rc->NumObjectIDsInScope(), 0);
+  ASSERT_EQ(out.size(), 4);
+}
+
 // Tests that the ref counts are properly integrated into the local
 // object memory store.
 TEST(MemoryStoreIntegrationTest, TestSimple) {