diff --git a/BUILD.bazel b/BUILD.bazel index 4d31f64f6..41ab0a35c 100644 --- a/BUILD.bazel +++ b/BUILD.bazel @@ -289,6 +289,7 @@ cc_library( "src/ray/object_manager/plasma/store_runner.cc", ], hdrs = [ + "src/ray/object_manager/common.h", "src/ray/object_manager/plasma/eviction_policy.h", "src/ray/object_manager/plasma/external_store.h", "src/ray/object_manager/plasma/plasma_allocator.h", diff --git a/python/ray/_private/services.py b/python/ray/_private/services.py index 41c8d22c2..426965385 100644 --- a/python/ray/_private/services.py +++ b/python/ray/_private/services.py @@ -1,3 +1,4 @@ +import base64 import collections import errno import io @@ -69,6 +70,21 @@ ProcessInfo = collections.namedtuple("ProcessInfo", [ ]) +def serialize_config(config): + config_pairs = [] + for key, value in config.items(): + if isinstance(value, str): + value = value.encode("utf-8") + if isinstance(value, bytes): + value = base64.b64encode(value).decode("utf-8") + config_pairs.append((key, value)) + config_str = ";".join(["{},{}".format(*kv) for kv in config_pairs]) + assert " " not in config_str, ( + "Config parameters currently do not support " + "spaces:", config_str) + return config_str + + class ConsolePopen(subprocess.Popen): if sys.platform == "win32": @@ -1121,7 +1137,7 @@ def start_gcs_server(redis_address, """ gcs_ip_address, gcs_port = redis_address.split(":") redis_password = redis_password or "" - config_str = ",".join(["{},{}".format(*kv) for kv in config.items()]) + config_str = serialize_config(config) if gcs_server_port is None: gcs_server_port = 0 @@ -1176,7 +1192,6 @@ def start_raylet(redis_address, socket_to_use=None, head_node=False, start_initial_python_workers_for_first_job=False, - object_spilling_config=None, code_search_path=None): """Start a raylet, which is a combined local scheduler and object manager. @@ -1223,7 +1238,7 @@ def start_raylet(redis_address, # The caller must provide a node manager port so that we can correctly # populate the command to start a worker. assert node_manager_port is not None and node_manager_port != 0 - config_str = ",".join(["{},{}".format(*kv) for kv in config.items()]) + config_str = serialize_config(config) if use_valgrind and use_profiler: raise ValueError("Cannot use valgrind and profiler at the same time.") @@ -1315,10 +1330,6 @@ def start_raylet(redis_address, if load_code_from_local: start_worker_command += ["--load-code-from-local"] - if object_spilling_config: - start_worker_command.append( - f"--object-spilling-config={json.dumps(object_spilling_config)}") - # Create agent command agent_command = [ sys.executable, diff --git a/python/ray/node.py b/python/ray/node.py index d645b70cd..7c4b6d00d 100644 --- a/python/ray/node.py +++ b/python/ray/node.py @@ -734,7 +734,6 @@ class Node: head_node=self.head, start_initial_python_workers_for_first_job=self._ray_params. start_initial_python_workers_for_first_job, - object_spilling_config=self._ray_params.object_spilling_config, code_search_path=self._ray_params.code_search_path) assert ray_constants.PROCESS_TYPE_RAYLET not in self.all_processes self.all_processes[ray_constants.PROCESS_TYPE_RAYLET] = [process_info] diff --git a/python/ray/parameter.py b/python/ray/parameter.py index 24193dcdb..6c9114603 100644 --- a/python/ray/parameter.py +++ b/python/ray/parameter.py @@ -1,3 +1,4 @@ +import json import logging import os @@ -148,7 +149,6 @@ class RayParams: metrics_agent_port=None, metrics_export_port=None, lru_evict=False, - object_spilling_config=None, code_search_path=None): self.object_ref_seed = object_ref_seed self.redis_address = redis_address @@ -191,10 +191,9 @@ class RayParams: self.metrics_export_port = metrics_export_port self.start_initial_python_workers_for_first_job = ( start_initial_python_workers_for_first_job) - self._system_config = _system_config + self._system_config = _system_config or {} self._lru_evict = lru_evict self._enable_object_reconstruction = enable_object_reconstruction - self.object_spilling_config = object_spilling_config self._check_usage() self.code_search_path = code_search_path if code_search_path is None: @@ -322,8 +321,10 @@ class RayParams: "serialization. Upgrade numpy if using with ray.") # Make sure object spilling configuration is applicable. - object_spilling_config = self.object_spilling_config or {} + object_spilling_config = self._system_config.get( + "object_spilling_config", {}) if object_spilling_config: + object_spilling_config = json.loads(object_spilling_config) from ray import external_storage # Validate external storage usage. external_storage.setup_external_storage(object_spilling_config) diff --git a/python/ray/tests/test_object_spilling.py b/python/ray/tests/test_object_spilling.py index de8237395..9950cf2c5 100644 --- a/python/ray/tests/test_object_spilling.py +++ b/python/ray/tests/test_object_spilling.py @@ -26,14 +26,16 @@ smart_open_object_spilling_config = { @pytest.fixture( - scope="module", + scope="function", params=[ file_system_object_spilling_config, # TODO(sang): Add a mock dependency to test S3. # smart_open_object_spilling_config, ]) -def object_spilling_config(request): - yield request.param +def object_spilling_config(request, tmpdir): + if request.param["type"] == "filesystem": + request.param["params"]["directory_path"] = str(tmpdir) + yield json.dumps(request.param) @pytest.mark.skip("This test is for local benchmark.") @@ -48,10 +50,10 @@ def test_sample_benchmark(object_spilling_config, shutdown_only): # Limit our object store to 200 MiB of memory. ray.init( object_store_memory=object_store_limit, - _object_spilling_config=object_spilling_config, _system_config={ "object_store_full_max_retries": 0, "max_io_workers": max_io_workers, + "object_spilling_config": object_spilling_config, }) arr = np.random.rand(object_size) replay_buffer = [] @@ -91,18 +93,26 @@ def test_invalid_config_raises_exception(shutdown_only): # it starts processes when invalid object spilling # config is given. with pytest.raises(ValueError): - ray.init(_object_spilling_config={"type": "abc"}) + ray.init(_system_config={ + "object_spilling_config": json.dumps({ + "type": "abc" + }), + }) with pytest.raises(Exception): copied_config = copy.deepcopy(file_system_object_spilling_config) # Add invalid params to the config. copied_config["params"].update({"random_arg": "abc"}) - ray.init(_object_spilling_config=copied_config) + ray.init(_system_config={ + "object_spilling_config": json.dumps(copied_config), + }) with pytest.raises(ValueError): copied_config = copy.deepcopy(file_system_object_spilling_config) copied_config["params"].update({"directory_path": "not_exist_path"}) - ray.init(_object_spilling_config=copied_config) + ray.init(_system_config={ + "object_spilling_config": json.dumps(copied_config), + }) @pytest.mark.skipif( @@ -111,10 +121,11 @@ def test_spill_objects_manually(object_spilling_config, shutdown_only): # Limit our object store to 75 MiB of memory. ray.init( object_store_memory=75 * 1024 * 1024, - _object_spilling_config=object_spilling_config, _system_config={ "object_store_full_max_retries": 0, + "automatic_object_spilling_enabled": False, "max_io_workers": 4, + "object_spilling_config": object_spilling_config, }) arr = np.random.rand(1024 * 1024) # 8 MB data replay_buffer = [] @@ -161,10 +172,11 @@ def test_spill_objects_manually_from_workers(object_spilling_config, # Limit our object store to 100 MiB of memory. ray.init( object_store_memory=100 * 1024 * 1024, - _object_spilling_config=object_spilling_config, _system_config={ "object_store_full_max_retries": 0, + "automatic_object_spilling_enabled": False, "max_io_workers": 4, + "object_spilling_config": object_spilling_config, }) @ray.remote @@ -190,10 +202,11 @@ def test_spill_objects_manually_with_workers(object_spilling_config, # Limit our object store to 75 MiB of memory. ray.init( object_store_memory=100 * 1024 * 1024, - _object_spilling_config=object_spilling_config, _system_config={ "object_store_full_max_retries": 0, + "automatic_object_spilling_enabled": False, "max_io_workers": 4, + "object_spilling_config": object_spilling_config, }) arrays = [np.random.rand(100 * 1024) for _ in range(50)] objects = [ray.put(arr) for arr in arrays] @@ -210,23 +223,27 @@ def test_spill_objects_manually_with_workers(object_spilling_config, @pytest.mark.skipif( platform.system() == "Windows", reason="Failing on Windows.") -def test_spill_remote_object(object_spilling_config, ray_start_cluster): - cluster = ray_start_cluster - # # Head node. - cluster.add_node( - num_cpus=0, - object_store_memory=75 * 1024 * 1024, - object_spilling_config=object_spilling_config, - _system_config={ - "object_store_full_max_retries": 0, +@pytest.mark.parametrize( + "ray_start_cluster_head", [{ + "num_cpus": 0, + "object_store_memory": 75 * 1024 * 1024, + "_system_config": { + "automatic_object_spilling_enabled": True, + "object_store_full_max_retries": 4, + "object_store_full_initial_delay_ms": 100, "max_io_workers": 4, - }) - # Worker nodes. - cluster.add_node( - object_store_memory=75 * 1024 * 1024, - object_spilling_config=object_spilling_config) - cluster.wait_for_nodes() - ray.init(address=cluster.address) + "object_spilling_config": json.dumps({ + "type": "filesystem", + "params": { + "directory_path": "/tmp" + } + }), + }, + }], + indirect=True) +def test_spill_remote_object(ray_start_cluster_head): + cluster = ray_start_cluster_head + cluster.add_node(object_store_memory=75 * 1024 * 1024) @ray.remote def put(): @@ -240,11 +257,7 @@ def test_spill_remote_object(object_spilling_config, ray_start_cluster): copy = np.copy(ray.get(ref)) # Evict local copy. ray.put(np.random.rand(5 * 1024 * 1024)) # 40 MB data - # Remote copy should not fit. - with pytest.raises(ray.exceptions.RayTaskError): - ray.get(put.remote()) - # Spill 1 object. The second should now fit. - ray.experimental.force_spill_objects([ref]) + # Remote copy should cause first remote object to get spilled. ray.get(put.remote()) sample = ray.get(ref) @@ -258,17 +271,17 @@ def test_spill_remote_object(object_spilling_config, ray_start_cluster): ray.get(depends.remote(ref)) -@pytest.mark.skip(reason="Not implemented yet.") -def test_spill_objects_automatically(shutdown_only): +def test_spill_objects_automatically(object_spilling_config, shutdown_only): # Limit our object store to 75 MiB of memory. ray.init( object_store_memory=75 * 1024 * 1024, - _system_config=json.dumps({ + _system_config={ "max_io_workers": 4, - "object_store_full_max_retries": 2, - "object_store_full_initial_delay_ms": 10, - "auto_object_spilling": True, - })) + "automatic_object_spilling_enabled": True, + "object_store_full_max_retries": 4, + "object_store_full_initial_delay_ms": 100, + "object_spilling_config": object_spilling_config, + }) arr = np.random.rand(1024 * 1024) # 8 MB data replay_buffer = [] @@ -291,5 +304,37 @@ def test_spill_objects_automatically(shutdown_only): assert np.array_equal(sample, arr) +@pytest.mark.skipif( + platform.system() == "Windows", reason="Failing on Windows.") +def test_spill_during_get(object_spilling_config, shutdown_only): + ray.init( + num_cpus=4, + object_store_memory=100 * 1024 * 1024, + _system_config={ + "automatic_object_spilling_enabled": True, + # This test will deadlock if only one IO worker is allowed because + # the IO worker will try to restore an object, but this requires + # another object to be spilled, which also requires an IO worker. + "max_io_workers": 2, + "object_spilling_config": object_spilling_config, + }, + ) + + @ray.remote + def f(): + return np.zeros(10 * 1024 * 1024) + + ids = [] + for i in range(10): + x = f.remote() + print(i, x) + ids.append(x) + + # Concurrent gets, which require restoring from external storage, while + # objects are being created. + for x in ids: + print(ray.get(x).shape) + + if __name__ == "__main__": sys.exit(pytest.main(["-sv", __file__])) diff --git a/python/ray/worker.py b/python/ray/worker.py index 0f94aac97..b6be97973 100644 --- a/python/ray/worker.py +++ b/python/ray/worker.py @@ -510,7 +510,6 @@ def init( _load_code_from_local=False, _lru_evict=False, _metrics_export_port=None, - _object_spilling_config=None, _system_config=None): """ Connect to an existing Ray cluster or start one and connect to it. @@ -609,8 +608,6 @@ def init( _metrics_export_port(int): Port number Ray exposes system metrics through a Prometheus endpoint. It is currently under active development, and the API is subject to change. - _object_spilling_config (str): The configuration json string for object - spilling I/O worker. _system_config (dict): Configuration for overriding RayConfig defaults. For testing purposes ONLY. @@ -727,8 +724,7 @@ def init( _system_config=_system_config, lru_evict=_lru_evict, enable_object_reconstruction=_enable_object_reconstruction, - metrics_export_port=_metrics_export_port, - object_spilling_config=_object_spilling_config) + metrics_export_port=_metrics_export_port) # Start the Ray processes. We set shutdown_at_exit=False because we # shutdown the node in the ray.shutdown call that happens in the atexit # handler. We still spawn a reaper process in case the atexit handler diff --git a/python/ray/workers/default_worker.py b/python/ray/workers/default_worker.py index d2534e841..3273fea32 100644 --- a/python/ray/workers/default_worker.py +++ b/python/ray/workers/default_worker.py @@ -1,4 +1,5 @@ import argparse +import base64 import json import time import sys @@ -125,20 +126,13 @@ if __name__ == "__main__": if mode == ray.IO_WORKER_MODE: from ray import external_storage if args.object_spilling_config: - object_spilling_config = json.loads(args.object_spilling_config) + object_spilling_config = base64.b64decode( + args.object_spilling_config) + object_spilling_config = json.loads(object_spilling_config) else: object_spilling_config = {} external_storage.setup_external_storage(object_spilling_config) - system_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): - system_config[config_list[i]] = config_list[i + 1] - i += 2 - raylet_ip_address = args.raylet_ip_address if raylet_ip_address is None: raylet_ip_address = args.node_ip_address @@ -161,7 +155,6 @@ if __name__ == "__main__": temp_dir=args.temp_dir, load_code_from_local=args.load_code_from_local, metrics_agent_port=args.metrics_agent_port, - _system_config=system_config, ) node = ray.node.Node( diff --git a/src/ray/common/ray_config_def.h b/src/ray/common/ray_config_def.h index a6ce69808..a2fd66228 100644 --- a/src/ray/common/ray_config_def.h +++ b/src/ray/common/ray_config_def.h @@ -314,9 +314,6 @@ RAY_CONFIG(bool, ownership_based_object_directory_enabled, false) // The interval where metrics are exported in milliseconds. RAY_CONFIG(uint64_t, metrics_report_interval_ms, 10000) -/// The maximum number of I/O worker that raylet starts. -RAY_CONFIG(int, max_io_workers, 1) - /// Enable the task timeline. If this is enabled, certain events such as task /// execution are profiled and sent to the GCS. RAY_CONFIG(bool, enable_timeline, true) @@ -324,3 +321,15 @@ RAY_CONFIG(bool, enable_timeline, true) /// The maximum number of pending placement group entries that are reported to monitor to /// autoscale the cluster. RAY_CONFIG(int64_t, max_placement_group_load_report_size, 100) + +/* Configuration parameters for object spilling. */ +/// JSON configuration that describes the external storage. This is passed to +/// Python IO workers to determine how to store/restore an object to/from +/// external storage. +RAY_CONFIG(std::string, object_spilling_config, "") +/// Whether to enable automatic object spilling. If enabled, then +/// Ray will choose objects to spill when the object store is out of +/// memory. +RAY_CONFIG(bool, automatic_object_spilling_enabled, true) +/// The maximum number of I/O worker that raylet starts. +RAY_CONFIG(int, max_io_workers, 1) diff --git a/src/ray/common/status.cc b/src/ray/common/status.cc index b3e2eba50..0c02accbe 100644 --- a/src/ray/common/status.cc +++ b/src/ray/common/status.cc @@ -56,6 +56,7 @@ namespace ray { #define STATUS_CODE_OBJECT_NOT_FOUND "ObjectNotFound" #define STATUS_CODE_OBJECT_STORE_ALREADY_SEALED "ObjectAlreadySealed" #define STATUS_CODE_OBJECT_STORE_FULL "ObjectStoreFull" +#define STATUS_CODE_TRANSIENT_OBJECT_STORE_FULL "TransientObjectStoreFull" Status::Status(StatusCode code, const std::string &msg) { assert(code != StatusCode::OK); @@ -98,6 +99,7 @@ std::string Status::CodeAsString() const { {StatusCode::ObjectNotFound, STATUS_CODE_OBJECT_NOT_FOUND}, {StatusCode::ObjectAlreadySealed, STATUS_CODE_OBJECT_STORE_ALREADY_SEALED}, {StatusCode::ObjectStoreFull, STATUS_CODE_OBJECT_STORE_FULL}, + {StatusCode::TransientObjectStoreFull, STATUS_CODE_TRANSIENT_OBJECT_STORE_FULL}, }; if (!code_to_str.count(code())) { diff --git a/src/ray/common/status.h b/src/ray/common/status.h index d8fd18400..2526d8915 100644 --- a/src/ray/common/status.h +++ b/src/ray/common/status.h @@ -99,6 +99,7 @@ enum class StatusCode : char { ObjectNotFound = 22, ObjectAlreadySealed = 23, ObjectStoreFull = 24, + TransientObjectStoreFull = 25, }; #if defined(__clang__) @@ -194,6 +195,10 @@ class RAY_EXPORT Status { return Status(StatusCode::ObjectStoreFull, msg); } + static Status TransientObjectStoreFull(const std::string &msg) { + return Status(StatusCode::TransientObjectStoreFull, msg); + } + // Returns true iff the status indicates success. bool ok() const { return (state_ == NULL); } @@ -220,6 +225,9 @@ class RAY_EXPORT Status { bool IsObjectNotFound() const { return code() == StatusCode::ObjectNotFound; } bool IsObjectAlreadySealed() const { return code() == StatusCode::ObjectAlreadySealed; } bool IsObjectStoreFull() const { return code() == StatusCode::ObjectStoreFull; } + bool IsTransientObjectStoreFull() const { + return code() == StatusCode::TransientObjectStoreFull; + } // Return a string representation of this status suitable for printing. // Returns the string "OK" for success. diff --git a/src/ray/core_worker/store_provider/plasma_store_provider.cc b/src/ray/core_worker/store_provider/plasma_store_provider.cc index f26647aa3..4bf15a2a7 100644 --- a/src/ray/core_worker/store_provider/plasma_store_provider.cc +++ b/src/ray/core_worker/store_provider/plasma_store_provider.cc @@ -127,6 +127,23 @@ Status CoreWorkerPlasmaStoreProvider::Create(const std::shared_ptr &meta "in the cluster." << "\n---\n"; } + } else if (plasma_status.IsTransientObjectStoreFull()) { + std::ostringstream message; + message << "Failed to put object " << object_id << " in object store because it " + << "is currently full, but space is being made through object spilling. " + "Object size is " + << data_size << " bytes."; + // The object store is full, but space is being made. Try again soon. + std::this_thread::sleep_for(std::chrono::milliseconds(10)); + delay = RayConfig::instance().object_store_full_initial_delay_ms(); + // Don't count these retries towards the total count. Also, reset the + // retry count since there may be space soon so we are not out of memory + // yet. + // NOTE(swang): We do this because the plasma store cannot guarantee that + // there will be enough space for the object on a future retry when there + // are concurrent clients trying to create objects. + retries = 0; + should_retry = true; } else if (plasma_status.IsObjectExists()) { RAY_LOG(WARNING) << "Trying to put an object that already existed in plasma: " << object_id << "."; diff --git a/src/ray/gcs/gcs_server/gcs_server_main.cc b/src/ray/gcs/gcs_server/gcs_server_main.cc index 85c3ef743..6c3879ca2 100644 --- a/src/ray/gcs/gcs_server/gcs_server_main.cc +++ b/src/ray/gcs/gcs_server/gcs_server_main.cc @@ -54,7 +54,7 @@ int main(int argc, char *argv[]) { std::string config_value; while (std::getline(config_string, config_name, ',')) { - RAY_CHECK(std::getline(config_string, config_value, ',')); + RAY_CHECK(std::getline(config_string, config_value, ';')); config_map[config_name] = config_value; } diff --git a/src/ray/object_manager/common.h b/src/ray/object_manager/common.h new file mode 100644 index 000000000..d48f72c96 --- /dev/null +++ b/src/ray/object_manager/common.h @@ -0,0 +1,11 @@ + +#pragma once + +namespace ray { + +/// A callback to asynchronously spill objects when space is needed. The +/// callback returns the amount of space still needed after the spilling is +/// complete. +using SpillObjectsCallback = std::function; + +} // namespace ray diff --git a/src/ray/object_manager/object_manager.cc b/src/ray/object_manager/object_manager.cc index c3df404e5..88a2af01a 100644 --- a/src/ray/object_manager/object_manager.cc +++ b/src/ray/object_manager/object_manager.cc @@ -26,14 +26,16 @@ namespace object_manager_protocol = ray::object_manager::protocol; namespace ray { -ObjectStoreRunner::ObjectStoreRunner(const ObjectManagerConfig &config) { +ObjectStoreRunner::ObjectStoreRunner(const ObjectManagerConfig &config, + SpillObjectsCallback spill_objects_callback) { if (config.object_store_memory > 0) { plasma::plasma_store_runner.reset(new plasma::PlasmaStoreRunner( config.store_socket_name, config.object_store_memory, config.huge_pages, config.plasma_directory, "")); // Initialize object store. store_thread_ = - std::thread(&plasma::PlasmaStoreRunner::Start, plasma::plasma_store_runner.get()); + std::thread(&plasma::PlasmaStoreRunner::Start, plasma::plasma_store_runner.get(), + spill_objects_callback); // Sleep for sometime until the store is working. This can suppress some // connection warnings. std::this_thread::sleep_for(std::chrono::microseconds(500)); @@ -51,11 +53,12 @@ ObjectStoreRunner::~ObjectStoreRunner() { ObjectManager::ObjectManager(asio::io_service &main_service, const NodeID &self_node_id, const ObjectManagerConfig &config, std::shared_ptr object_directory, - RestoreSpilledObjectCallback restore_spilled_object) + RestoreSpilledObjectCallback restore_spilled_object, + SpillObjectsCallback spill_objects_callback) : self_node_id_(self_node_id), config_(config), object_directory_(std::move(object_directory)), - object_store_internal_(config), + object_store_internal_(config, spill_objects_callback), buffer_pool_(config_.store_socket_name, config_.object_chunk_size), rpc_work_(rpc_service_), gen_(std::chrono::high_resolution_clock::now().time_since_epoch().count()), diff --git a/src/ray/object_manager/object_manager.h b/src/ray/object_manager/object_manager.h index b972d2771..e716c6624 100644 --- a/src/ray/object_manager/object_manager.h +++ b/src/ray/object_manager/object_manager.h @@ -81,7 +81,8 @@ struct LocalObjectInfo { class ObjectStoreRunner { public: - ObjectStoreRunner(const ObjectManagerConfig &config); + ObjectStoreRunner(const ObjectManagerConfig &config, + SpillObjectsCallback spill_objects_callback); ~ObjectStoreRunner(); private: @@ -190,7 +191,8 @@ class ObjectManager : public ObjectManagerInterface, explicit ObjectManager(boost::asio::io_service &main_service, const NodeID &self_node_id, const ObjectManagerConfig &config, std::shared_ptr object_directory, - RestoreSpilledObjectCallback restore_spilled_object); + RestoreSpilledObjectCallback restore_spilled_object, + SpillObjectsCallback spill_objects_callback = nullptr); ~ObjectManager(); diff --git a/src/ray/object_manager/plasma/eviction_policy.cc b/src/ray/object_manager/plasma/eviction_policy.cc index 4c053e82e..7640a6cea 100644 --- a/src/ray/object_manager/plasma/eviction_policy.cc +++ b/src/ray/object_manager/plasma/eviction_policy.cc @@ -120,7 +120,7 @@ bool EvictionPolicy::EnforcePerClientQuota(Client* client, int64_t size, bool is void EvictionPolicy::ClientDisconnected(Client* client) {} -bool EvictionPolicy::RequireSpace(int64_t size, std::vector* objects_to_evict) { +int64_t EvictionPolicy::RequireSpace(int64_t size, std::vector* objects_to_evict) { // Check if there is enough space to create the object. int64_t required_space = PlasmaAllocator::Allocated() + size - PlasmaAllocator::GetFootprintLimit(); @@ -135,7 +135,7 @@ bool EvictionPolicy::RequireSpace(int64_t size, std::vector* objects_t << objects_to_evict->size() << " objects to free up " << num_bytes_evicted << " bytes. The number of bytes in use (before " << "this eviction) is " << PlasmaAllocator::Allocated() << "."; - return num_bytes_evicted >= required_space && num_bytes_evicted > 0; + return required_space - num_bytes_evicted; } void EvictionPolicy::BeginObjectAccess(const ObjectID& object_id) { diff --git a/src/ray/object_manager/plasma/eviction_policy.h b/src/ray/object_manager/plasma/eviction_policy.h index 6abb1a6ec..0f68008b1 100644 --- a/src/ray/object_manager/plasma/eviction_policy.h +++ b/src/ray/object_manager/plasma/eviction_policy.h @@ -152,8 +152,9 @@ class EvictionPolicy { /// metadata. /// \param objects_to_evict The object IDs that were chosen for eviction will /// be stored into this vector. - /// \return True if enough space can be freed and false otherwise. - virtual bool RequireSpace(int64_t size, std::vector* objects_to_evict); + /// \return The number of bytes of space that is still needed, if + /// any. If negative, then the required space has been made. + virtual int64_t RequireSpace(int64_t size, std::vector* objects_to_evict); /// This method will be called whenever an unused object in the Plasma store /// starts to be used. When this method is called, the eviction policy will diff --git a/src/ray/object_manager/plasma/plasma.fbs b/src/ray/object_manager/plasma/plasma.fbs index e7e4fdcd8..456f73cfc 100644 --- a/src/ray/object_manager/plasma/plasma.fbs +++ b/src/ray/object_manager/plasma/plasma.fbs @@ -81,6 +81,11 @@ enum PlasmaError:int { ObjectNonexistent, // Trying to create an object but there isn't enough space in the store. OutOfMemory, + // Trying to create an object but there isn't enough space in the store. + // However, objects are currently being spilled to make enough space. The + // client should try again soon, and there will be enough space (assuming the + // space is not taken by another client). + TransientOutOfMemory, // Trying to delete an object but it's not sealed. ObjectNotSealed, // Trying to delete an object but it's in use. diff --git a/src/ray/object_manager/plasma/protocol.cc b/src/ray/object_manager/plasma/protocol.cc index c8f0d539e..12abda11f 100644 --- a/src/ray/object_manager/plasma/protocol.cc +++ b/src/ray/object_manager/plasma/protocol.cc @@ -128,6 +128,8 @@ Status PlasmaErrorStatus(fb::PlasmaError plasma_error) { return Status::ObjectNotFound("object does not exist in the plasma store"); case fb::PlasmaError::OutOfMemory: return Status::ObjectStoreFull("object does not fit in the plasma store"); + case fb::PlasmaError::TransientOutOfMemory: + return Status::TransientObjectStoreFull("object does not currently fit in the plasma store, try again soon"); default: RAY_LOG(FATAL) << "unknown plasma error code " << static_cast(plasma_error); } diff --git a/src/ray/object_manager/plasma/store.cc b/src/ray/object_manager/plasma/store.cc index 4156bb1d0..c68fac0c5 100644 --- a/src/ray/object_manager/plasma/store.cc +++ b/src/ray/object_manager/plasma/store.cc @@ -105,13 +105,15 @@ GetRequest::GetRequest(boost::asio::io_service& io_context, const std::shared_pt PlasmaStore::PlasmaStore(boost::asio::io_service &main_service, std::string directory, bool hugepages_enabled, const std::string& socket_name, - std::shared_ptr external_store) + std::shared_ptr external_store, + ray::SpillObjectsCallback spill_objects_callback) : io_context_(main_service), socket_name_(socket_name), acceptor_(main_service, ParseUrlEndpoint(socket_name)), socket_(main_service), eviction_policy_(&store_info_, PlasmaAllocator::GetFootprintLimit()), - external_store_(external_store) { + external_store_(external_store), + spill_objects_callback_(spill_objects_callback) { store_info_.directory = directory; store_info_.hugepages_enabled = hugepages_enabled; #ifdef PLASMA_CUDA @@ -159,13 +161,14 @@ void PlasmaStore::AddToClientObjectIds(const ObjectID& object_id, ObjectTableEnt // Allocate memory uint8_t* PlasmaStore::AllocateMemory(size_t size, bool evict_if_full, MEMFD_TYPE* fd, int64_t* map_size, ptrdiff_t* offset, const std::shared_ptr &client, - bool is_create) { + bool is_create, PlasmaError *error) { // First free up space from the client's LRU queue if quota enforcement is on. if (evict_if_full) { std::vector client_objects_to_evict; bool quota_ok = eviction_policy_.EnforcePerClientQuota(client.get(), size, is_create, &client_objects_to_evict); if (!quota_ok) { + *error = PlasmaError::OutOfMemory; return nullptr; } EvictObjects(client_objects_to_evict); @@ -186,15 +189,39 @@ uint8_t* PlasmaStore::AllocateMemory(size_t size, bool evict_if_full, MEMFD_TYPE // If we manage to allocate the memory, return the pointer. If we cannot // allocate the space, but we are also not allowed to evict anything to // make more space, return an error to the client. + *error = PlasmaError::OutOfMemory; break; } // Tell the eviction policy how much space we need to create this object. std::vector objects_to_evict; - bool success = eviction_policy_.RequireSpace(size, &objects_to_evict); + int64_t space_needed = eviction_policy_.RequireSpace(size, &objects_to_evict); EvictObjects(objects_to_evict); - // Return an error to the client if not enough space could be freed to - // create the object. - if (!success) { + // More space is still needed. Try to spill objects to external storage to + // make room. + if (space_needed > 0) { + if (spill_objects_callback_) { + // Object spilling is asynchronous so that we do not block the plasma + // store thread. Therefore the client must try again, even if enough + // space will be made after the spill is complete. + // TODO(swang): Only respond to the client with OutOfMemory if we could not + // make enough space through spilling. If we could make enough space, + // respond to the plasma client once spilling is complete. + space_needed = spill_objects_callback_(space_needed); + } + if (space_needed > 0) { + // There is still not enough space, even once all evictable objects + // were evicted and all pending object spills have finished. The + // client may choose to try again, or throw an OutOfMemory error to + // the application immediately. + *error = PlasmaError::OutOfMemory; + } else { + // Once all pending object spills have finished, there should be + // enough space for this allocation. Return a transient error to the + // client so that they try again soon. + *error = PlasmaError::TransientOutOfMemory; + } + // Return an error to the client if not enough space could be freed to + // create the object. break; } } @@ -202,6 +229,7 @@ uint8_t* PlasmaStore::AllocateMemory(size_t size, bool evict_if_full, MEMFD_TYPE if (pointer != nullptr) { GetMallocMapinfo(pointer, fd, map_size, offset); RAY_CHECK(*fd != INVALID_FD); + *error = PlasmaError::OK; } return pointer; } @@ -250,14 +278,17 @@ PlasmaError PlasmaStore::CreateObject(const ObjectID& object_id, auto total_size = data_size + metadata_size; if (device_num == 0) { + PlasmaError error = PlasmaError::OK; pointer = - AllocateMemory(total_size, evict_if_full, &fd, &map_size, &offset, client, true); + AllocateMemory(total_size, evict_if_full, &fd, &map_size, &offset, client, true, &error); if (!pointer) { - RAY_LOG(ERROR) << "Not enough memory to create the object " << object_id.Hex() - << ", data_size=" << data_size - << ", metadata_size=" << metadata_size - << ", will send a reply of PlasmaError::OutOfMemory"; - return PlasmaError::OutOfMemory; + if (error == PlasmaError::OutOfMemory) { + RAY_LOG(ERROR) << "Not enough memory to create the object " << object_id.Hex() + << ", data_size=" << data_size + << ", metadata_size=" << metadata_size + << ", will send a reply of PlasmaError::OutOfMemory"; + } + return error; } } else { #ifdef PLASMA_CUDA @@ -474,9 +505,11 @@ void PlasmaStore::ProcessGetRequest(const std::shared_ptr &client, // Make sure the object pointer is not already allocated RAY_CHECK(!entry->pointer); + PlasmaError error = PlasmaError::OK; entry->pointer = AllocateMemory(entry->data_size + entry->metadata_size, /*evict=*/true, - &entry->fd, &entry->map_size, &entry->offset, client, false); + &entry->fd, &entry->map_size, &entry->offset, client, + false, &error); if (entry->pointer) { entry->state = ObjectState::PLASMA_CREATED; entry->create_time = std::time(nullptr); @@ -602,8 +635,8 @@ ObjectStatus PlasmaStore::ContainsObject(const ObjectID& object_id) { void PlasmaStore::SealObjects(const std::vector& object_ids) { std::vector infos; - RAY_LOG(DEBUG) << "sealing " << object_ids.size() << " objects"; for (size_t i = 0; i < object_ids.size(); ++i) { + RAY_LOG(DEBUG) << "sealing object " << object_ids[i]; ObjectInfoT object_info; auto entry = GetObjectTableEntry(&store_info_, object_ids[i]); RAY_CHECK(entry != nullptr); diff --git a/src/ray/object_manager/plasma/store.h b/src/ray/object_manager/plasma/store.h index 29bf3cf04..191449706 100644 --- a/src/ray/object_manager/plasma/store.h +++ b/src/ray/object_manager/plasma/store.h @@ -25,6 +25,7 @@ #include #include "ray/common/status.h" +#include "ray/object_manager/common.h" #include "ray/object_manager/format/object_manager_generated.h" #include "ray/object_manager/notification/object_store_notification_manager.h" #include "ray/object_manager/plasma/common.h" @@ -52,7 +53,8 @@ class PlasmaStore { // TODO: PascalCase PlasmaStore methods. PlasmaStore(boost::asio::io_service &main_service, std::string directory, bool hugepages_enabled, const std::string& socket_name, - std::shared_ptr external_store); + std::shared_ptr external_store, + ray::SpillObjectsCallback spill_objects_callback); ~PlasmaStore(); @@ -94,6 +96,9 @@ class PlasmaStore { /// - PlasmaError::OutOfMemory, if the store is out of memory and /// cannot create the object. In this case, the client should not call /// plasma_release. + /// - PlasmaError::TransientOutOfMemory, if the store is temporarily out of + /// memory but there may be space soon to create the object. In this + /// case, the client should not call plasma_release. PlasmaError CreateObject(const ObjectID& object_id, const NodeID& owner_raylet_id, const std::string& owner_ip_address, int owner_port, const WorkerID& owner_worker_id, bool evict_if_full, @@ -222,7 +227,8 @@ class PlasmaStore { void EraseFromObjectTable(const ObjectID& object_id); uint8_t* AllocateMemory(size_t size, bool evict_if_full, MEMFD_TYPE* fd, int64_t* map_size, - ptrdiff_t* offset, const std::shared_ptr &client, bool is_create); + ptrdiff_t* offset, const std::shared_ptr &client, bool is_create, + PlasmaError *error); #ifdef PLASMA_CUDA Status AllocateCudaMemory(int device_num, int64_t size, uint8_t** out_pointer, std::shared_ptr* out_ipc_handle); @@ -262,6 +268,11 @@ class PlasmaStore { arrow::cuda::CudaDeviceManager* manager_; #endif std::shared_ptr notification_listener_; + /// A callback to asynchronously spill objects when space is needed. The + /// callback returns the amount of space still needed after the spilling is + /// complete. + ray::SpillObjectsCallback spill_objects_callback_; + }; } // namespace plasma diff --git a/src/ray/object_manager/plasma/store_runner.cc b/src/ray/object_manager/plasma/store_runner.cc index 8eb9c287d..f6c8a9abe 100644 --- a/src/ray/object_manager/plasma/store_runner.cc +++ b/src/ray/object_manager/plasma/store_runner.cc @@ -75,7 +75,7 @@ PlasmaStoreRunner::PlasmaStoreRunner(std::string socket_name, int64_t system_mem plasma_directory_ = plasma_directory; } -void PlasmaStoreRunner::Start() { +void PlasmaStoreRunner::Start(ray::SpillObjectsCallback spill_objects_callback) { // Get external store std::shared_ptr external_store{nullptr}; if (!external_store_endpoint_.empty()) { @@ -94,7 +94,8 @@ void PlasmaStoreRunner::Start() { { absl::MutexLock lock(&store_runner_mutex_); store_.reset(new PlasmaStore(main_service_, plasma_directory_, hugepages_enabled_, - socket_name_, external_store)); + socket_name_, external_store, + spill_objects_callback)); plasma_config = store_->GetPlasmaStoreInfo(); // We are using a single memory-mapped file by mallocing and freeing a single diff --git a/src/ray/object_manager/plasma/store_runner.h b/src/ray/object_manager/plasma/store_runner.h index c55535415..e9c953dfc 100644 --- a/src/ray/object_manager/plasma/store_runner.h +++ b/src/ray/object_manager/plasma/store_runner.h @@ -15,7 +15,7 @@ class PlasmaStoreRunner { PlasmaStoreRunner(std::string socket_name, int64_t system_memory, bool hugepages_enabled, std::string plasma_directory, const std::string external_store_endpoint); - void Start(); + void Start(ray::SpillObjectsCallback spill_objects_callback = nullptr); void Stop(); void SetNotificationListener( const std::shared_ptr ¬ification_listener) { diff --git a/src/ray/raylet/local_object_manager.cc b/src/ray/raylet/local_object_manager.cc index fce7536b3..2fc8e63dc 100644 --- a/src/ray/raylet/local_object_manager.cc +++ b/src/ray/raylet/local_object_manager.cc @@ -20,6 +20,7 @@ namespace raylet { void LocalObjectManager::PinObjects(const std::vector &object_ids, std::vector> &&objects) { + absl::MutexLock lock(&mutex_); for (size_t i = 0; i < object_ids.size(); i++) { const auto &object_id = object_ids[i]; auto &object = objects[i]; @@ -56,8 +57,11 @@ void LocalObjectManager::WaitForObjectFree(const rpc::Address &owner_address, } void LocalObjectManager::ReleaseFreedObject(const ObjectID &object_id) { - RAY_LOG(DEBUG) << "Unpinning object " << object_id; - pinned_objects_.erase(object_id); + { + absl::MutexLock lock(&mutex_); + RAY_LOG(DEBUG) << "Unpinning object " << object_id; + pinned_objects_.erase(object_id); + } // Try to evict all copies of the object from the cluster. if (free_objects_period_ms_ >= 0) { @@ -85,36 +89,114 @@ void LocalObjectManager::FlushFreeObjectsIfNeeded(int64_t now_ms) { } } +int64_t LocalObjectManager::SpillObjectsOfSize(int64_t num_bytes_required) { + if (RayConfig::instance().object_spilling_config().empty() || + !RayConfig::instance().automatic_object_spilling_enabled()) { + return num_bytes_required; + } + + absl::MutexLock lock(&mutex_); + + RAY_LOG(INFO) << "Choosing objects to spill of total size " << num_bytes_required; + int64_t num_bytes_to_spill = 0; + auto it = pinned_objects_.begin(); + std::vector objects_to_spill; + while (num_bytes_to_spill < num_bytes_required && it != pinned_objects_.end()) { + num_bytes_to_spill += it->second->GetSize(); + objects_to_spill.push_back(it->first); + it++; + } + if (!objects_to_spill.empty()) { + RAY_LOG(ERROR) << "Spilling objects of total size " << num_bytes_to_spill; + auto start_time = current_time_ms(); + SpillObjectsInternal( + objects_to_spill, [num_bytes_to_spill, start_time](const Status &status) { + if (!status.ok()) { + RAY_LOG(ERROR) << "Error spilling objects " << status.ToString(); + } else { + RAY_LOG(INFO) << "Spilled " << num_bytes_to_spill << " in " + << (current_time_ms() - start_time) << "ms"; + } + }); + } + // We do not track a mapping between objects that need to be created to + // objects that are being spilled, so we just subtract the total number of + // bytes that are currently being spilled from the amount of space + // requested. If the space is claimed by another client, this client may + // need to request space again. + num_bytes_required -= num_bytes_pending_spill_; + return num_bytes_required; +} + void LocalObjectManager::SpillObjects(const std::vector &object_ids, std::function callback) { + absl::MutexLock lock(&mutex_); + SpillObjectsInternal(object_ids, callback); +} + +void LocalObjectManager::SpillObjectsInternal( + const std::vector &object_ids, + std::function callback) { + std::vector objects_to_spill; + // Filter for the objects that can be spilled. for (const auto &id : object_ids) { - // We should not spill an object that we are not the primary copy for. - if (pinned_objects_.count(id) == 0) { - callback( - Status::Invalid("Requested spill for object that is not marked as " - "the primary copy.")); + // We should not spill an object that we are not the primary copy for, or + // objects that are already being spilled. + if (pinned_objects_.count(id) == 0 && objects_pending_spill_.count(id) == 0) { + if (callback) { + callback( + Status::Invalid("Requested spill for object that is not marked as " + "the primary copy.")); + } + return; + } + + // Add objects that we are the primary copy for, and that we are not + // already spilling. + auto it = pinned_objects_.find(id); + if (it != pinned_objects_.end()) { + RAY_LOG(DEBUG) << "Spilling object " << id; + objects_to_spill.push_back(id); + num_bytes_pending_spill_ += it->second->GetSize(); + objects_pending_spill_[id] = std::move(it->second); + pinned_objects_.erase(it); } } + if (objects_to_spill.empty()) { + if (callback) { + callback(Status::Invalid("All objects are already being spilled.")); + } + return; + } + io_worker_pool_.PopIOWorker( - [this, object_ids, callback](std::shared_ptr io_worker) { + [this, objects_to_spill, callback](std::shared_ptr io_worker) { rpc::SpillObjectsRequest request; - for (const auto &object_id : object_ids) { + for (const auto &object_id : objects_to_spill) { RAY_LOG(DEBUG) << "Sending spill request for object " << object_id; request.add_object_ids_to_spill(object_id.Binary()); } io_worker->rpc_client()->SpillObjects( - request, [this, object_ids, callback, io_worker]( + request, [this, objects_to_spill, callback, io_worker]( const ray::Status &status, const rpc::SpillObjectsReply &r) { io_worker_pool_.PushIOWorker(io_worker); + absl::MutexLock lock(&mutex_); if (!status.ok()) { + for (const auto &object_id : objects_to_spill) { + auto it = objects_pending_spill_.find(object_id); + RAY_CHECK(it != objects_pending_spill_.end()); + pinned_objects_.emplace(object_id, std::move(it->second)); + objects_pending_spill_.erase(it); + } + RAY_LOG(ERROR) << "Failed to send object spilling request: " << status.ToString(); if (callback) { callback(status); } } else { - AddSpilledUrls(object_ids, r, callback); + AddSpilledUrls(objects_to_spill, r, callback); } }); }); @@ -134,12 +216,13 @@ void LocalObjectManager::AddSpilledUrls( RAY_CHECK_OK(object_info_accessor_.AsyncAddSpilledUrl( object_id, object_url, [this, object_id, callback, num_remaining](Status status) { RAY_CHECK_OK(status); + absl::MutexLock lock(&mutex_); // Unpin the object. - // NOTE(swang): Due to a race condition, the object may not be in - // the map yet. In that case, the owner will respond to the - // WaitForObjectEvictionRequest and we will unpin the object - // then. - pinned_objects_.erase(object_id); + auto it = objects_pending_spill_.find(object_id); + RAY_CHECK(it != objects_pending_spill_.end()); + num_bytes_pending_spill_ -= it->second->GetSize(); + objects_pending_spill_.erase(it); + (*num_remaining)--; if (*num_remaining == 0 && callback) { callback(status); @@ -153,18 +236,21 @@ void LocalObjectManager::AsyncRestoreSpilledObject( std::function callback) { RAY_LOG(DEBUG) << "Restoring spilled object " << object_id << " from URL " << object_url; - io_worker_pool_.PopIOWorker([this, object_url, + io_worker_pool_.PopIOWorker([this, object_id, object_url, callback](std::shared_ptr io_worker) { RAY_LOG(DEBUG) << "Sending restore spilled object request"; rpc::RestoreSpilledObjectsRequest request; request.add_spilled_objects_url(std::move(object_url)); io_worker->rpc_client()->RestoreSpilledObjects( - request, [this, callback, io_worker](const ray::Status &status, - const rpc::RestoreSpilledObjectsReply &r) { + request, + [this, object_id, callback, io_worker](const ray::Status &status, + const rpc::RestoreSpilledObjectsReply &r) { io_worker_pool_.PushIOWorker(io_worker); if (!status.ok()) { RAY_LOG(ERROR) << "Failed to send restore spilled object request: " << status.ToString(); + } else { + RAY_LOG(DEBUG) << "Restored object " << object_id; } if (callback) { callback(status); diff --git a/src/ray/raylet/local_object_manager.h b/src/ray/raylet/local_object_manager.h index 719c03d15..153bb6b8c 100644 --- a/src/ray/raylet/local_object_manager.h +++ b/src/ray/raylet/local_object_manager.h @@ -62,6 +62,14 @@ class LocalObjectManager { void WaitForObjectFree(const rpc::Address &owner_address, const std::vector &object_ids); + /// Asynchronously spill objects whose total size adds up to at least the + /// specified number of bytes. + /// + /// \param num_bytes_to_spill The total number of bytes to spill. + /// \return The number of bytes of space still required after the spill is + /// complete. + int64_t SpillObjectsOfSize(int64_t num_bytes_to_spill); + /// Spill objects to external storage. /// /// \param objects_ids_to_spill The objects to be spilled. @@ -83,6 +91,11 @@ class LocalObjectManager { void FlushFreeObjectsIfNeeded(int64_t now_ms); private: + /// Internal helper method for spilling objects. + void SpillObjectsInternal(const std::vector &objects_ids, + std::function callback) + EXCLUSIVE_LOCKS_REQUIRED(mutex_); + /// Release an object that has been freed by its owner. void ReleaseFreedObject(const ObjectID &object_id); @@ -116,7 +129,14 @@ class LocalObjectManager { std::function &)> on_objects_freed_; // Objects that are pinned on this node. - absl::flat_hash_map> pinned_objects_; + absl::flat_hash_map> pinned_objects_ + GUARDED_BY(mutex_); + + // Objects that were pinned on this node but that are being spilled. + // These objects will be released once spilling is complete and the URL is + // written to the object directory. + absl::flat_hash_map> objects_pending_spill_ + GUARDED_BY(mutex_); /// The time that we last sent a FreeObjects request to other nodes for /// objects that have gone out of scope in the application. @@ -127,6 +147,14 @@ class LocalObjectManager { /// free_objects_batch_size, or if objects have been in the cache for longer /// than the config's free_objects_period, whichever occurs first. std::vector objects_to_free_; + + /// The total size of the objects that are currently being + /// spilled from this node, in bytes. + size_t num_bytes_pending_spill_ GUARDED_BY(mutex_) = 0; + + /// This class is accessed by both the raylet and plasma store threads. The + /// mutex protects private members that relate to object spilling. + mutable absl::Mutex mutex_; }; }; // namespace raylet diff --git a/src/ray/raylet/main.cc b/src/ray/raylet/main.cc index 816f52017..a20843612 100644 --- a/src/ray/raylet/main.cc +++ b/src/ray/raylet/main.cc @@ -127,7 +127,7 @@ int main(int argc, char *argv[]) { std::string config_value; while (std::getline(config_string, config_name, ',')) { - RAY_CHECK(std::getline(config_string, config_value, ',')); + RAY_CHECK(std::getline(config_string, config_value, ';')); // TODO(rkn): The line below could throw an exception. What should we do about this? raylet_config[config_name] = config_value; } diff --git a/src/ray/raylet/raylet.cc b/src/ray/raylet/raylet.cc index 80f8aab27..93ec573ca 100644 --- a/src/ray/raylet/raylet.cc +++ b/src/ray/raylet/raylet.cc @@ -75,6 +75,10 @@ Raylet::Raylet(boost::asio::io_service &main_service, const std::string &socket_ std::function callback) { node_manager_.GetLocalObjectManager().AsyncRestoreSpilledObject( object_id, spilled_url, callback); + }, + [this](int64_t num_bytes_required) { + return node_manager_.GetLocalObjectManager().SpillObjectsOfSize( + num_bytes_required); }), node_manager_(main_service, self_node_id_, node_manager_config, object_manager_, gcs_client_, object_directory_), diff --git a/src/ray/raylet/test/local_object_manager_test.cc b/src/ray/raylet/test/local_object_manager_test.cc index 9bdb2f9fc..c13f1cc22 100644 --- a/src/ray/raylet/test/local_object_manager_test.cc +++ b/src/ray/raylet/test/local_object_manager_test.cc @@ -68,8 +68,10 @@ class MockIOWorkerClient : public rpc::CoreWorkerClientInterface { } auto callback = callbacks.front(); auto reply = rpc::SpillObjectsReply(); - for (const auto &url : urls) { - reply.add_spilled_objects_url(url); + if (status.ok()) { + for (const auto &url : urls) { + reply.add_spilled_objects_url(url); + } } callback(status, reply); callbacks.pop_front(); @@ -128,10 +130,20 @@ class MockObjectInfoAccessor : public gcs::ObjectInfoAccessor { Status AsyncAddSpilledUrl(const ObjectID &object_id, const std::string &spilled_url, const gcs::StatusCallback &callback) { object_urls[object_id] = spilled_url; - callback(Status()); + callbacks.push_back(callback); return Status(); } + bool ReplyAsyncAddSpilledUrl() { + if (callbacks.size() == 0) { + return false; + } + auto callback = callbacks.front(); + callback(Status::OK()); + callbacks.pop_front(); + return true; + } + MOCK_METHOD3(AsyncRemoveLocation, Status(const ObjectID &object_id, const NodeID &node_id, const gcs::StatusCallback &callback)); @@ -149,6 +161,28 @@ class MockObjectInfoAccessor : public gcs::ObjectInfoAccessor { MOCK_METHOD1(IsObjectUnsubscribed, bool(const ObjectID &object_id)); std::unordered_map object_urls; + std::list callbacks; +}; + +class MockObjectBuffer : public Buffer { + public: + MockObjectBuffer(size_t size, ObjectID object_id, + std::shared_ptr> unpins) + : size_(size), id_(object_id), unpins_(unpins) {} + + MOCK_CONST_METHOD0(Data, uint8_t *()); + + size_t Size() const { return size_; } + + MOCK_CONST_METHOD0(OwnsData, bool()); + + MOCK_CONST_METHOD0(IsPlasmaBuffer, bool()); + + ~MockObjectBuffer() { (*unpins_)[id_]++; } + + size_t size_; + ObjectID id_; + std::shared_ptr> unpins_; }; class LocalObjectManagerTest : public ::testing::Test { @@ -162,7 +196,10 @@ class LocalObjectManagerTest : public ::testing::Test { for (const auto &object_id : object_ids) { freed.insert(object_id); } - }) {} + }), + unpins(std::make_shared>()) { + RayConfig::instance().initialize({{"object_spilling_config", "mock_config"}}); + } size_t free_objects_batch_size = 3; std::shared_ptr owner_client; @@ -172,6 +209,9 @@ class LocalObjectManagerTest : public ::testing::Test { LocalObjectManager manager; std::unordered_set freed; + // This hashmap is incremented when objects are unpinned by destroying their + // unique_ptr. + std::shared_ptr> unpins; }; TEST_F(LocalObjectManagerTest, TestPin) { @@ -221,11 +261,9 @@ TEST_F(LocalObjectManagerTest, TestExplicitSpill) { for (size_t i = 0; i < free_objects_batch_size; i++) { ObjectID object_id = ObjectID::FromRandom(); object_ids.push_back(object_id); - std::string meta = std::to_string(static_cast(rpc::ErrorType::OBJECT_IN_PLASMA)); - auto metadata = const_cast(reinterpret_cast(meta.data())); - auto meta_buffer = std::make_shared(metadata, meta.size()); + auto data_buffer = std::make_shared(0, object_id, unpins); std::unique_ptr object( - new RayObject(nullptr, meta_buffer, std::vector())); + new RayObject(data_buffer, nullptr, std::vector())); objects.push_back(std::move(object)); } manager.PinObjects(object_ids, std::move(objects)); @@ -236,6 +274,9 @@ TEST_F(LocalObjectManagerTest, TestExplicitSpill) { num_times_fired++; }); ASSERT_EQ(num_times_fired, 0); + for (const auto &id : object_ids) { + ASSERT_EQ((*unpins)[id], 0); + } EXPECT_CALL(worker_pool, PushIOWorker(_)); std::vector urls; @@ -243,10 +284,165 @@ TEST_F(LocalObjectManagerTest, TestExplicitSpill) { urls.push_back("url" + std::to_string(i)); } ASSERT_TRUE(worker_pool.io_worker_client->ReplySpillObjects(urls)); + for (size_t i = 0; i < object_ids.size(); i++) { + ASSERT_TRUE(object_table.ReplyAsyncAddSpilledUrl()); + } ASSERT_EQ(num_times_fired, 1); for (size_t i = 0; i < object_ids.size(); i++) { ASSERT_EQ(object_table.object_urls[object_ids[i]], urls[i]); } + for (const auto &id : object_ids) { + ASSERT_EQ((*unpins)[id], 1); + } +} + +TEST_F(LocalObjectManagerTest, TestDuplicateSpill) { + rpc::Address owner_address; + owner_address.set_worker_id(WorkerID::FromRandom().Binary()); + + std::vector object_ids; + std::vector> objects; + + for (size_t i = 0; i < free_objects_batch_size; i++) { + ObjectID object_id = ObjectID::FromRandom(); + object_ids.push_back(object_id); + auto data_buffer = std::make_shared(0, object_id, unpins); + std::unique_ptr object( + new RayObject(data_buffer, nullptr, std::vector())); + objects.push_back(std::move(object)); + } + manager.PinObjects(object_ids, std::move(objects)); + manager.WaitForObjectFree(owner_address, object_ids); + + int num_times_fired = 0; + manager.SpillObjects(object_ids, [&](const Status &status) mutable { + ASSERT_TRUE(status.ok()); + num_times_fired++; + }); + // Spill the same objects again. The callback should only be fired once + // total. + manager.SpillObjects(object_ids, + [&](const Status &status) mutable { ASSERT_TRUE(!status.ok()); }); + ASSERT_EQ(num_times_fired, 0); + for (const auto &id : object_ids) { + ASSERT_EQ((*unpins)[id], 0); + } + + std::vector urls; + for (size_t i = 0; i < object_ids.size(); i++) { + urls.push_back("url" + std::to_string(i)); + } + EXPECT_CALL(worker_pool, PushIOWorker(_)); + ASSERT_TRUE(worker_pool.io_worker_client->ReplySpillObjects(urls)); + for (size_t i = 0; i < object_ids.size(); i++) { + ASSERT_TRUE(object_table.ReplyAsyncAddSpilledUrl()); + } + ASSERT_EQ(num_times_fired, 1); + for (size_t i = 0; i < object_ids.size(); i++) { + ASSERT_EQ(object_table.object_urls[object_ids[i]], urls[i]); + } + ASSERT_FALSE(worker_pool.io_worker_client->ReplySpillObjects(urls)); + for (const auto &id : object_ids) { + ASSERT_EQ((*unpins)[id], 1); + } +} + +TEST_F(LocalObjectManagerTest, TestSpillObjectsOfSize) { + rpc::Address owner_address; + owner_address.set_worker_id(WorkerID::FromRandom().Binary()); + + std::vector object_ids; + std::vector> objects; + int64_t total_size = 0; + int64_t object_size = 1000; + + for (size_t i = 0; i < 3; i++) { + ObjectID object_id = ObjectID::FromRandom(); + object_ids.push_back(object_id); + auto data_buffer = std::make_shared(object_size, object_id, unpins); + total_size += object_size; + std::unique_ptr object( + new RayObject(data_buffer, nullptr, std::vector())); + objects.push_back(std::move(object)); + } + manager.PinObjects(object_ids, std::move(objects)); + + int64_t num_bytes_required = manager.SpillObjectsOfSize(total_size / 2); + ASSERT_EQ(num_bytes_required, -object_size / 2); + for (const auto &id : object_ids) { + ASSERT_EQ((*unpins)[id], 0); + } + + // Check that this returns the total number of bytes currently being spilled. + num_bytes_required = manager.SpillObjectsOfSize(0); + ASSERT_EQ(num_bytes_required, -2 * object_size); + + // Check that half the objects get spilled and the URLs get added to the + // global object directory. + std::vector urls; + for (size_t i = 0; i < object_ids.size() / 2 + 1; i++) { + urls.push_back("url" + std::to_string(i)); + } + EXPECT_CALL(worker_pool, PushIOWorker(_)); + // Objects should get freed even though we didn't wait for the owner's notice + // to evict. + ASSERT_TRUE(worker_pool.io_worker_client->ReplySpillObjects(urls)); + for (size_t i = 0; i < urls.size(); i++) { + ASSERT_TRUE(object_table.ReplyAsyncAddSpilledUrl()); + } + ASSERT_EQ(object_table.object_urls.size(), object_ids.size() / 2 + 1); + for (auto &object_url : object_table.object_urls) { + auto it = std::find(urls.begin(), urls.end(), object_url.second); + ASSERT_TRUE(it != urls.end()); + ASSERT_EQ((*unpins)[object_url.first], 1); + } + + // Check that this returns the total number of bytes currently being spilled. + num_bytes_required = manager.SpillObjectsOfSize(0); + ASSERT_EQ(num_bytes_required, 0); +} + +TEST_F(LocalObjectManagerTest, TestSpillError) { + // Check that we can spill an object again if there was a transient error + // during the first attempt. + rpc::Address owner_address; + owner_address.set_worker_id(WorkerID::FromRandom().Binary()); + + ObjectID object_id = ObjectID::FromRandom(); + auto data_buffer = std::make_shared(0, object_id, unpins); + std::unique_ptr object( + new RayObject(std::move(data_buffer), nullptr, std::vector())); + + std::vector> objects; + objects.push_back(std::move(object)); + manager.PinObjects({object_id}, std::move(objects)); + + int num_times_fired = 0; + manager.SpillObjects({object_id}, [&](const Status &status) mutable { + ASSERT_FALSE(status.ok()); + num_times_fired++; + }); + + // Return an error from the IO worker during spill. + EXPECT_CALL(worker_pool, PushIOWorker(_)); + ASSERT_TRUE( + worker_pool.io_worker_client->ReplySpillObjects({}, Status::IOError("error"))); + ASSERT_FALSE(object_table.ReplyAsyncAddSpilledUrl()); + ASSERT_EQ(num_times_fired, 1); + ASSERT_EQ((*unpins)[object_id], 0); + + // Try to spill the same object again. + manager.SpillObjects({object_id}, [&](const Status &status) mutable { + ASSERT_TRUE(status.ok()); + num_times_fired++; + }); + std::string url = "url"; + EXPECT_CALL(worker_pool, PushIOWorker(_)); + ASSERT_TRUE(worker_pool.io_worker_client->ReplySpillObjects({url})); + ASSERT_TRUE(object_table.ReplyAsyncAddSpilledUrl()); + ASSERT_EQ(num_times_fired, 2); + ASSERT_EQ(object_table.object_urls[object_id], url); + ASSERT_EQ((*unpins)[object_id], 1); } } // namespace raylet diff --git a/src/ray/raylet/worker_pool.cc b/src/ray/raylet/worker_pool.cc index ccf0303a6..16e9296cc 100644 --- a/src/ray/raylet/worker_pool.cc +++ b/src/ray/raylet/worker_pool.cc @@ -320,6 +320,14 @@ Process WorkerPool::StartWorkerProcess( } } + if (worker_type == rpc::WorkerType::IO_WORKER) { + RAY_CHECK(!RayConfig::instance().object_spilling_config().empty()); + RAY_LOG(INFO) << "Adding object spill config " + << RayConfig::instance().object_spilling_config(); + worker_command_args.push_back("--object-spilling-config=" + + RayConfig::instance().object_spilling_config()); + } + ProcessEnvironment env; if (RayConfig::instance().enable_multi_tenancy() && job_config) { env.insert(job_config->worker_env().begin(), job_config->worker_env().end());