diff --git a/python/ray/_raylet.pyx b/python/ray/_raylet.pyx index 2d5b51933..4f450ef2b 100644 --- a/python/ray/_raylet.pyx +++ b/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 diff --git a/python/ray/includes/ray_config.pxd b/python/ray/includes/ray_config.pxd index fa412f527..2d61c9820 100644 --- a/python/ray/includes/ray_config.pxd +++ b/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) diff --git a/python/ray/services.py b/python/ray/services.py index d3852186f..fe59035fe 100644 --- a/python/ray/services.py +++ b/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) diff --git a/python/ray/worker.py b/python/ray/worker.py index bb11cd327..908dd10cb 100644 --- a/python/ray/worker.py +++ b/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, diff --git a/python/ray/workers/default_worker.py b/python/ray/workers/default_worker.py index 2f178a3ad..93f5f22b0 100644 --- a/python/ray/workers/default_worker.py +++ b/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() diff --git a/src/ray/common/ray_config.h b/src/ray/common/ray_config.h index 544868552..0f5685ae3 100644 --- a/src/ray/common/ray_config.h +++ b/src/ray/common/ray_config.h @@ -43,21 +43,15 @@ class RayConfig { } void initialize(const std::unordered_map &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 diff --git a/src/ray/core_worker/core_worker.cc b/src/ray/core_worker/core_worker.cc index 287dcabde..51dc442e6 100644 --- a/src/ray/core_worker/core_worker.cc +++ b/src/ray/core_worker/core_worker.cc @@ -1,11 +1,11 @@ +#include "ray/core_worker/core_worker.h" + #include #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 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 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 &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; } diff --git a/src/ray/core_worker/core_worker.h b/src/ray/core_worker/core_worker.h index 2746a9f7c..5479cef1b 100644 --- a/src/ray/core_worker/core_worker.h +++ b/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> *args, std::vector *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 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_; diff --git a/src/ray/core_worker/reference_count.cc b/src/ray/core_worker/reference_count.cc index 6e8b4fcd0..b0b484a79 100644 --- a/src/ray/core_worker/reference_count.cc +++ b/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 *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 *deleted) { auto entry = object_id_refs_.find(object_id); diff --git a/src/ray/core_worker/reference_count.h b/src/ray/core_worker/reference_count.h index 96e5ff759..bb63d0297 100644 --- a/src/ray/core_worker/reference_count.h +++ b/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> 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 *deleted) + LOCKS_EXCLUDED(mutex_); + /// Returns the total number of ObjectIDs currently in scope. size_t NumObjectIDsInScope() const LOCKS_EXCLUDED(mutex_); diff --git a/src/ray/core_worker/reference_count_test.cc b/src/ray/core_worker/reference_count_test.cc index 004d74efb..f177a85b9 100644 --- a/src/ray/core_worker/reference_count_test.cc +++ b/src/ray/core_worker/reference_count_test.cc @@ -1,8 +1,9 @@ +#include "ray/core_worker/reference_count.h" + #include #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 out; + ObjectID id1 = ObjectID::FromRandom(); + ObjectID id2 = ObjectID::FromRandom(); + ObjectID id3 = ObjectID::FromRandom(); + ObjectID id4 = ObjectID::FromRandom(); + + std::shared_ptr> deps2 = + std::make_shared>(); + deps2->push_back(id3); + deps2->push_back(id4); + rc->AddOwnedObject(id2, TaskID::Nil(), rpc::Address(), deps2); + + std::shared_ptr> deps1 = + std::make_shared>(); + 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) {