[Object Spilling] Fusion small objects (#12087)

src/ray/common/ray_config_def.h

@@ -332,9 +332,17 @@ RAY_CONFIG(int64_t, max_placement_group_load_report_size, 100)
 /// 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)
+
+/// Ray's object spilling fuses small objects into a single file before flushing them
+/// to optimize the performance.
+/// The minimum object size that can be spilled by each spill operation. 100 MB by
+/// default. This value is not recommended to set beyond --object-store-memory.
+RAY_CONFIG(int64_t, min_spilling_size, 100 * 1024 * 1024)

src/ray/core_worker/core_worker.cc

@@ -2346,12 +2346,19 @@ void CoreWorker::HandleRestoreSpilledObjects(
     const rpc::RestoreSpilledObjectsRequest &request,
     rpc::RestoreSpilledObjectsReply *reply, rpc::SendReplyCallback send_reply_callback) {
   if (options_.restore_spilled_objects != nullptr) {
+    // Get a list of object ids.
+    std::vector<ObjectID> object_ids_to_restore;
+    object_ids_to_restore.reserve(request.object_ids_to_restore_size());
+    for (const auto &id_binary : request.object_ids_to_restore()) {
+      object_ids_to_restore.push_back(ObjectID::FromBinary(id_binary));
+    }
+    // Get a list of spilled_object_urls.
     std::vector<std::string> spilled_objects_url;
     spilled_objects_url.reserve(request.spilled_objects_url_size());
     for (const auto &url : request.spilled_objects_url()) {
       spilled_objects_url.push_back(url);
     }
-    options_.restore_spilled_objects(spilled_objects_url);
+    options_.restore_spilled_objects(object_ids_to_restore, spilled_objects_url);
     send_reply_callback(Status::OK(), nullptr, nullptr);
   } else {
     send_reply_callback(

src/ray/core_worker/core_worker.h

@@ -138,7 +138,8 @@ struct CoreWorkerOptions {
   /// Application-language callback to spill objects to external storage.
   std::function<std::vector<std::string>(const std::vector<ObjectID> &)> spill_objects;
   /// Application-language callback to restore objects from external storage.
-  std::function<void(const std::vector<std::string> &)> restore_spilled_objects;
+  std::function<void(const std::vector<ObjectID> &, const std::vector<std::string> &)>
+      restore_spilled_objects;
   /// Language worker callback to get the current call stack.
   std::function<void(std::string *)> get_lang_stack;
   // Function that tries to interrupt the currently running Python thread.

src/ray/object_manager/common.h

@@ -5,10 +5,15 @@
 
 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<int64_t(int64_t num_bytes_required)>;
+/// A callback to asynchronously spill objects when space is needed.
+/// The callback tries to spill objects as much as num_bytes_to_spill and returns
+/// the amount of space needed after the spilling is complete.
+/// The returned value is calculated based off of min_bytes_to_spill. That says,
+/// although it fails to spill num_bytes_to_spill, as long as it spills more than
+/// min_bytes_to_spill, it will return the value that is less than 0 (meaning we
+/// don't need any more additional space).
+using SpillObjectsCallback =
+    std::function<int64_t(int64_t num_bytes_to_spill, int64_t min_bytes_to_spill)>;
 
 /// A callback to call when space has been released.
 using SpaceReleasedCallback = std::function<void()>;

src/ray/object_manager/plasma/store.cc

@@ -216,13 +216,18 @@ uint8_t *PlasmaStore::AllocateMemory(size_t size, bool evict_if_full, MEMFD_TYPE
     // make room.
     if (space_needed > 0) {
       if (spill_objects_callback_) {
+        // If the space needed is too small, we'd like to bump up to the minimum spilling
+        // size. Cap the max size to be lower than the plasma store limit.
+        int64_t byte_to_spill =
+            std::min(PlasmaAllocator::GetFootprintLimit(),
+                     std::max(space_needed, RayConfig::instance().min_spilling_size()));
         // 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);
+        space_needed = spill_objects_callback_(byte_to_spill, space_needed);
       }
       if (space_needed > 0) {
         // There is still not enough space, even once all evictable objects

src/ray/protobuf/core_worker.proto

@@ -309,6 +309,9 @@ message SpillObjectsReply {
 message RestoreSpilledObjectsRequest {
   // The URLs of spilled objects.
   repeated string spilled_objects_url = 1;
+  // Object ids to restore. The order of object ids
+  // must be the same as spilled_objects_url.
+  repeated bytes object_ids_to_restore = 2;
 }
 
 message RestoreSpilledObjectsReply {

src/ray/raylet/local_object_manager.cc

@@ -89,32 +89,34 @@ void LocalObjectManager::FlushFreeObjectsIfNeeded(int64_t now_ms) {
   }
 }
 
-int64_t LocalObjectManager::SpillObjectsOfSize(int64_t num_bytes_required) {
+int64_t LocalObjectManager::SpillObjectsOfSize(int64_t num_bytes_to_spill,
+                                               int64_t min_bytes_to_spill) {
+  RAY_CHECK(num_bytes_to_spill >= min_bytes_to_spill);
   if (RayConfig::instance().object_spilling_config().empty() ||
       !RayConfig::instance().automatic_object_spilling_enabled()) {
-    return num_bytes_required;
+    return min_bytes_to_spill;
   }
 
   absl::MutexLock lock(&mutex_);
 
-  RAY_LOG(INFO) << "Choosing objects to spill of total size " << num_bytes_required;
-  int64_t num_bytes_to_spill = 0;
+  RAY_LOG(INFO) << "Choosing objects to spill of total size " << num_bytes_to_spill;
+  int64_t bytes_to_spill = 0;
   auto it = pinned_objects_.begin();
   std::vector<ObjectID> objects_to_spill;
-  while (num_bytes_to_spill < num_bytes_required && it != pinned_objects_.end()) {
-    num_bytes_to_spill += it->second->GetSize();
+  while (bytes_to_spill < num_bytes_to_spill && it != pinned_objects_.end()) {
+    bytes_to_spill += it->second->GetSize();
     objects_to_spill.push_back(it->first);
     it++;
   }
   if (!objects_to_spill.empty()) {
-    RAY_LOG(INFO) << "Spilling objects of total size " << num_bytes_to_spill;
+    RAY_LOG(INFO) << "Spilling objects of total size " << bytes_to_spill;
     auto start_time = current_time_ms();
     SpillObjectsInternal(
-        objects_to_spill, [num_bytes_to_spill, start_time](const Status &status) {
+        objects_to_spill, [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 "
+            RAY_LOG(INFO) << "Spilled " << bytes_to_spill << " in "
                           << (current_time_ms() - start_time) << "ms";
           }
         });
@@ -124,8 +126,7 @@ int64_t LocalObjectManager::SpillObjectsOfSize(int64_t num_bytes_required) {
   //  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;
+  return min_bytes_to_spill - num_bytes_pending_spill_;
 }
 
 void LocalObjectManager::SpillObjects(const std::vector<ObjectID> &object_ids,
@@ -169,7 +170,6 @@ void LocalObjectManager::SpillObjectsInternal(
     }
     return;
   }
-
   io_worker_pool_.PopSpillWorker(
       [this, objects_to_spill, callback](std::shared_ptr<WorkerInterface> io_worker) {
         rpc::SpillObjectsRequest request;
@@ -246,6 +246,7 @@ void LocalObjectManager::AsyncRestoreSpilledObject(
     RAY_LOG(DEBUG) << "Sending restore spilled object request";
     rpc::RestoreSpilledObjectsRequest request;
     request.add_spilled_objects_url(std::move(object_url));
+    request.add_object_ids_to_restore(object_id.Binary());
     io_worker->rpc_client()->RestoreSpilledObjects(
         request,
         [this, object_id, callback, io_worker](const ray::Status &status,

src/ray/raylet/local_object_manager.h

@@ -65,13 +65,22 @@ class LocalObjectManager {
   void WaitForObjectFree(const rpc::Address &owner_address,
                          const std::vector<ObjectID> &object_ids);
 
-  /// Asynchronously spill objects whose total size adds up to at least the
-  /// specified number of bytes.
+  /// Asynchronously spill objects when space is needed.
+  /// The callback tries to spill objects as much as num_bytes_to_spill and returns
+  /// the amount of space needed after the spilling is complete.
+  /// The returned value is calculated based off of min_bytes_to_spill. That says,
+  /// although it fails to spill num_bytes_to_spill, as long as it spills more than
+  /// min_bytes_to_spill, it will return the value that is less than 0 (meaning we
+  /// don't need any more additional space).
   ///
-  /// \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);
+  /// \param num_bytes_to_spill The total number of bytes to spill. The method tries to
+  /// spill bytes as much as this value.
+  /// \param min_bytes_to_spill The minimum bytes that
+  /// need to be spilled.
+  /// \return The number of bytes of space still required after the
+  /// spill is complete. This return the value is less than 0 if it satifies the
+  /// min_bytes_to_spill.
+  int64_t SpillObjectsOfSize(int64_t num_bytes_to_spill, int64_t min_bytes_to_spill);
 
   /// Spill objects to external storage.
   ///

src/ray/raylet/raylet.cc

@@ -76,9 +76,9 @@ Raylet::Raylet(boost::asio::io_service &main_service, const std::string &socket_
                         node_manager_.GetLocalObjectManager().AsyncRestoreSpilledObject(
                             object_id, spilled_url, callback);
                       },
-                      [this](int64_t num_bytes_required) {
+                      [this](int64_t num_bytes_to_spill, int64_t min_bytes_to_spill) {
                         return node_manager_.GetLocalObjectManager().SpillObjectsOfSize(
-                            num_bytes_required);
+                            num_bytes_to_spill, min_bytes_to_spill);
                       }),
       node_manager_(main_service, self_node_id_, node_manager_config, object_manager_,
                     gcs_client_, object_directory_,

src/ray/raylet/test/local_object_manager_test.cc

@@ -380,14 +380,14 @@ TEST_F(LocalObjectManagerTest, TestSpillObjectsOfSize) {
   }
   manager.PinObjects(object_ids, std::move(objects));
 
-  int64_t num_bytes_required = manager.SpillObjectsOfSize(total_size / 2);
+  int64_t num_bytes_required = manager.SpillObjectsOfSize(total_size / 2, 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);
+  num_bytes_required = manager.SpillObjectsOfSize(0, 0);
   ASSERT_EQ(num_bytes_required, -2 * object_size);
 
   // Check that half the objects get spilled and the URLs get added to the
@@ -411,7 +411,7 @@ TEST_F(LocalObjectManagerTest, TestSpillObjectsOfSize) {
   }
 
   // Check that this returns the total number of bytes currently being spilled.
-  num_bytes_required = manager.SpillObjectsOfSize(0);
+  num_bytes_required = manager.SpillObjectsOfSize(0, 0);
   ASSERT_EQ(num_bytes_required, 0);
   ASSERT_TRUE(num_callbacks_fired > 0);
 }
@@ -460,6 +460,52 @@ TEST_F(LocalObjectManagerTest, TestSpillError) {
   ASSERT_TRUE(num_callbacks_fired > 0);
 }
 
+TEST_F(LocalObjectManagerTest,
+       TestSpillObjectsOfSizeNumBytesToSpillHigherThanMinBytesToSpill) {
+  /// Test the case SpillObjectsOfSize(num_bytes_to_spill, min_bytes_to_spill
+  /// where num_bytes_to_spill > min_bytes_to_spill.
+  rpc::Address owner_address;
+  owner_address.set_worker_id(WorkerID::FromRandom().Binary());
+
+  std::vector<ObjectID> object_ids;
+  std::vector<std::unique_ptr<RayObject>> objects;
+  int64_t total_size = 0;
+  int64_t object_size = 1000;
+  size_t object_len = 3;
+
+  for (size_t i = 0; i < object_len; i++) {
+    ObjectID object_id = ObjectID::FromRandom();
+    object_ids.push_back(object_id);
+    auto data_buffer = std::make_shared<MockObjectBuffer>(object_size, object_id, unpins);
+    total_size += object_size;
+    std::unique_ptr<RayObject> object(
+        new RayObject(data_buffer, nullptr, std::vector<ObjectID>()));
+    objects.push_back(std::move(object));
+  }
+  manager.PinObjects(object_ids, std::move(objects));
+
+  // First test when num_bytes_to_spill > min_bytes to spill.
+  // It means that we cannot spill the num_bytes_required, but we at least spilled the
+  // required amount, which is the min_bytes_to_spill.
+  int64_t num_bytes_required = manager.SpillObjectsOfSize(8000, object_size);
+  // only min bytes to spill is considered.
+  ASSERT_TRUE(num_bytes_required <= 0);
+
+  // Make sure the spilling is done properly.
+  std::vector<std::string> urls;
+  for (size_t i = 0; i < object_ids.size(); i++) {
+    urls.push_back("url" + std::to_string(i));
+  }
+  EXPECT_CALL(worker_pool, PushSpillWorker(_));
+  ASSERT_TRUE(worker_pool.io_worker_client->ReplySpillObjects(urls));
+  for (size_t i = 0; i < object_ids.size(); i++) {
+    ASSERT_TRUE(object_table.ReplyAsyncAddSpilledUrl());
+  }
+  for (size_t i = 0; i < object_ids.size(); i++) {
+    ASSERT_EQ((*unpins).size(), object_len);
+  }
+}
+
 }  // namespace raylet
 
 }  // namespace ray