[Object spilling] Refactor raylet to add a local object manager class (#11647)

* Fix pytest...

* Release objects that have been spilled

* GCS object table interface refactor

* Add spilled URL to object location info

* refactor to include spilled URL in notifications

* improve tests

* Add spilled URL to object directory results

* Remove force restore call

* Merge spilled URL and location

* fix

* tmp

* refactor

* unit test skeleton

* unit testing

* unit test fixes

* cleanup

* cleanup

* update

* Separate pinning from waiting for object free, fixes pytest

* Update src/ray/raylet/local_object_manager.h

Co-authored-by: Eric Liang <ekhliang@gmail.com>

Co-authored-by: Tyler Westenbroek <westenbroekt@berkeley.edu>
Co-authored-by: Eric Liang <ekhliang@gmail.com>

BUILD.bazel

@@ -794,6 +794,18 @@ cc_test(
     ],
 )
 
+cc_test(
+    name = "local_object_manager_test",
+    srcs = [
+        "src/ray/raylet/test/local_object_manager_test.cc",
+    ],
+    copts = COPTS,
+    deps = [
+        ":raylet_lib",
+        "@com_google_googletest//:gtest_main",
+    ],
+)
+
 cc_test(
     name = "reconstruction_policy_test",
     srcs = ["src/ray/raylet/reconstruction_policy_test.cc"],

src/ray/raylet/local_object_manager.cc

@@ -0,0 +1,178 @@
+// Copyright 2017 The Ray Authors.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//  http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include "ray/raylet/local_object_manager.h"
+
+namespace ray {
+
+namespace raylet {
+
+void LocalObjectManager::PinObjects(const std::vector<ObjectID> &object_ids,
+                                    std::vector<std::unique_ptr<RayObject>> &&objects) {
+  for (size_t i = 0; i < object_ids.size(); i++) {
+    const auto &object_id = object_ids[i];
+    auto &object = objects[i];
+    if (object == nullptr) {
+      RAY_LOG(ERROR) << "Plasma object " << object_id
+                     << " was evicted before the raylet could pin it.";
+      continue;
+    }
+    RAY_LOG(DEBUG) << "Pinning object " << object_id;
+    pinned_objects_.emplace(object_id, std::move(object));
+  }
+}
+
+void LocalObjectManager::WaitForObjectFree(const rpc::Address &owner_address,
+                                           const std::vector<ObjectID> &object_ids) {
+  for (const auto &object_id : object_ids) {
+    // Send a long-running RPC request to the owner for each object. When we get a
+    // response or the RPC fails (due to the owner crashing), unpin the object.
+    // TODO(edoakes): we should be batching these requests instead of sending one per
+    // pinned object.
+    rpc::WaitForObjectEvictionRequest wait_request;
+    wait_request.set_object_id(object_id.Binary());
+    wait_request.set_intended_worker_id(owner_address.worker_id());
+    auto owner_client = owner_client_pool_.GetOrConnect(owner_address);
+    owner_client->WaitForObjectEviction(
+        wait_request,
+        [this, object_id](Status status, const rpc::WaitForObjectEvictionReply &reply) {
+          if (!status.ok()) {
+            RAY_LOG(WARNING) << "Worker failed. Unpinning object " << object_id;
+          }
+          ReleaseFreedObject(object_id);
+        });
+  }
+}
+
+void LocalObjectManager::ReleaseFreedObject(const ObjectID &object_id) {
+  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) {
+    objects_to_free_.push_back(object_id);
+  }
+  if (objects_to_free_.size() == free_objects_batch_size_ ||
+      free_objects_period_ms_ == 0) {
+    FlushFreeObjects();
+  }
+}
+
+void LocalObjectManager::FlushFreeObjects() {
+  if (!objects_to_free_.empty()) {
+    RAY_LOG(DEBUG) << "Freeing " << objects_to_free_.size() << " out-of-scope objects";
+    on_objects_freed_(objects_to_free_);
+    objects_to_free_.clear();
+  }
+  last_free_objects_at_ms_ = current_time_ms();
+}
+
+void LocalObjectManager::FlushFreeObjectsIfNeeded(int64_t now_ms) {
+  if (free_objects_period_ms_ > 0 &&
+      static_cast<int64_t>(now_ms - last_free_objects_at_ms_) > free_objects_period_ms_) {
+    FlushFreeObjects();
+  }
+}
+
+void LocalObjectManager::SpillObjects(const std::vector<ObjectID> &object_ids,
+                                      std::function<void(const ray::Status &)> callback) {
+  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."));
+    }
+  }
+
+  io_worker_pool_.PopIOWorker(
+      [this, object_ids, callback](std::shared_ptr<WorkerInterface> io_worker) {
+        rpc::SpillObjectsRequest request;
+        for (const auto &object_id : object_ids) {
+          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](
+                         const ray::Status &status, const rpc::SpillObjectsReply &r) {
+              io_worker_pool_.PushIOWorker(io_worker);
+              if (!status.ok()) {
+                RAY_LOG(ERROR) << "Failed to send object spilling request: "
+                               << status.ToString();
+                if (callback) {
+                  callback(status);
+                }
+              } else {
+                AddSpilledUrls(object_ids, r, callback);
+              }
+            });
+      });
+}
+
+void LocalObjectManager::AddSpilledUrls(
+    const std::vector<ObjectID> &object_ids, const rpc::SpillObjectsReply &worker_reply,
+    std::function<void(const ray::Status &)> callback) {
+  auto num_remaining = std::make_shared<size_t>(object_ids.size());
+  for (size_t i = 0; i < object_ids.size(); ++i) {
+    const ObjectID &object_id = object_ids[i];
+    const std::string &object_url = worker_reply.spilled_objects_url(i);
+    RAY_LOG(DEBUG) << "Object " << object_id << " spilled at " << object_url;
+    // Write to object directory. Wait for the write to finish before
+    // releasing the object to make sure that the spilled object can
+    // be retrieved by other raylets.
+    RAY_CHECK_OK(object_info_accessor_.AsyncAddSpilledUrl(
+        object_id, object_url, [this, object_id, callback, num_remaining](Status status) {
+          RAY_CHECK_OK(status);
+          // 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);
+          (*num_remaining)--;
+          if (*num_remaining == 0 && callback) {
+            callback(status);
+          }
+        }));
+  }
+}
+
+void LocalObjectManager::AsyncRestoreSpilledObject(
+    const ObjectID &object_id, const std::string &object_url,
+    std::function<void(const ray::Status &)> callback) {
+  RAY_LOG(DEBUG) << "Restoring spilled object " << object_id << " from URL "
+                 << object_url;
+  io_worker_pool_.PopIOWorker([this, object_url,
+                               callback](std::shared_ptr<WorkerInterface> 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) {
+          io_worker_pool_.PushIOWorker(io_worker);
+          if (!status.ok()) {
+            RAY_LOG(ERROR) << "Failed to send restore spilled object request: "
+                           << status.ToString();
+          }
+          if (callback) {
+            callback(status);
+          }
+        });
+  });
+}
+
+};  // namespace raylet
+
+};  // namespace ray

src/ray/raylet/local_object_manager.h

@@ -0,0 +1,134 @@
+// Copyright 2017 The Ray Authors.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//  http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#pragma once
+
+#include <google/protobuf/repeated_field.h>
+
+#include <functional>
+
+#include "ray/common/id.h"
+#include "ray/common/ray_object.h"
+#include "ray/gcs/accessor.h"
+#include "ray/raylet/worker_pool.h"
+#include "ray/rpc/worker/core_worker_client_pool.h"
+
+namespace ray {
+
+namespace raylet {
+
+/// This class implements memory management for primary objects, objects that
+/// have been freed, and objects that have been spilled.
+class LocalObjectManager {
+ public:
+  LocalObjectManager(size_t free_objects_batch_size, int64_t free_objects_period_ms,
+                     IOWorkerPoolInterface &io_worker_pool,
+                     gcs::ObjectInfoAccessor &object_info_accessor,
+                     rpc::CoreWorkerClientPool &owner_client_pool,
+                     std::function<void(const std::vector<ObjectID> &)> on_objects_freed)
+      : free_objects_period_ms_(free_objects_period_ms),
+        free_objects_batch_size_(free_objects_batch_size),
+        io_worker_pool_(io_worker_pool),
+        object_info_accessor_(object_info_accessor),
+        owner_client_pool_(owner_client_pool),
+        on_objects_freed_(on_objects_freed),
+        last_free_objects_at_ms_(current_time_ms()) {}
+
+  /// Pin objects.
+  ///
+  /// \param object_ids The objects to be pinned.
+  /// \param objects Pointers to the objects to be pinned. The pointer should
+  /// be kept in scope until the object can be released.
+  void PinObjects(const std::vector<ObjectID> &object_ids,
+                  std::vector<std::unique_ptr<RayObject>> &&objects);
+
+  /// Wait for the objects' owner to free the object.  The objects will be
+  /// released when the owner at the given address fails or replies that the
+  /// object can be evicted.
+  ///
+  /// \param owner_address The address of the owner of the objects.
+  /// \param object_ids The objects to be freed.
+  void WaitForObjectFree(const rpc::Address &owner_address,
+                         const std::vector<ObjectID> &object_ids);
+
+  /// Spill objects to external storage.
+  ///
+  /// \param objects_ids_to_spill The objects to be spilled.
+  /// \param callback A callback to call once the objects have been spilled, or
+  /// there is an error.
+  void SpillObjects(const std::vector<ObjectID> &objects_ids,
+                    std::function<void(const ray::Status &)> callback);
+
+  /// Restore a spilled object from external storage back into local memory.
+  ///
+  /// \param object_id The ID of the object to restore.
+  /// \param object_url The URL in external storage from which the object can be restored.
+  /// \param callback A callback to call when the restoration is done. Status
+  /// will contain the error during restoration, if any.
+  void AsyncRestoreSpilledObject(const ObjectID &object_id, const std::string &object_url,
+                                 std::function<void(const ray::Status &)> callback);
+
+  /// Try to clear any objects that have been freed.
+  void FlushFreeObjectsIfNeeded(int64_t now_ms);
+
+ private:
+  /// Release an object that has been freed by its owner.
+  void ReleaseFreedObject(const ObjectID &object_id);
+
+  /// Clear any freed objects. This will trigger the callback for freed
+  /// objects.
+  void FlushFreeObjects();
+
+  /// Add objects' spilled URLs to the global object directory. Call the
+  /// callback once all URLs have been added.
+  void AddSpilledUrls(const std::vector<ObjectID> &object_ids,
+                      const rpc::SpillObjectsReply &worker_reply,
+                      std::function<void(const ray::Status &)> callback);
+
+  /// The period between attempts to eagerly evict objects from plasma.
+  const int64_t free_objects_period_ms_;
+
+  /// The number of freed objects to accumulate before flushing.
+  const size_t free_objects_batch_size_;
+
+  /// A worker pool, used for spilling and restoring objects.
+  IOWorkerPoolInterface &io_worker_pool_;
+
+  /// A GCS client, used to update locations for spilled objects.
+  gcs::ObjectInfoAccessor &object_info_accessor_;
+
+  /// Cache of gRPC clients to owners of objects pinned on
+  /// this node.
+  rpc::CoreWorkerClientPool &owner_client_pool_;
+
+  /// A callback to call when an object has been freed.
+  std::function<void(const std::vector<ObjectID> &)> on_objects_freed_;
+
+  // Objects that are pinned on this node.
+  absl::flat_hash_map<ObjectID, std::unique_ptr<RayObject>> pinned_objects_;
+
+  /// The time that we last sent a FreeObjects request to other nodes for
+  /// objects that have gone out of scope in the application.
+  uint64_t last_free_objects_at_ms_ = 0;
+
+  /// Objects that are out of scope in the application and that should be freed
+  /// from plasma. The cache is flushed when it reaches the
+  /// 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<ObjectID> objects_to_free_;
+};
+
+};  // namespace raylet
+
+};  // namespace ray

src/ray/raylet/main.cc

@@ -219,8 +219,6 @@ int main(int argc, char *argv[]) {
             RayConfig::instance().raylet_heartbeat_timeout_milliseconds();
         node_manager_config.debug_dump_period_ms =
             RayConfig::instance().debug_dump_period_milliseconds();
-        node_manager_config.free_objects_period_ms =
-            RayConfig::instance().free_objects_period_milliseconds();
         node_manager_config.fair_queueing_enabled =
             RayConfig::instance().fair_queueing_enabled();
         node_manager_config.object_pinning_enabled =

src/ray/raylet/node_manager.cc

@@ -130,7 +130,6 @@ NodeManager::NodeManager(boost::asio::io_service &io_service, const NodeID &self
       heartbeat_timer_(io_service),
       heartbeat_period_(std::chrono::milliseconds(config.heartbeat_period_ms)),
       debug_dump_period_(config.debug_dump_period_ms),
-      free_objects_period_(config.free_objects_period_ms),
       fair_queueing_enabled_(config.fair_queueing_enabled),
       object_pinning_enabled_(config.object_pinning_enabled),
       temp_dir_(config.temp_dir),
@@ -162,6 +161,14 @@ NodeManager::NodeManager(boost::asio::io_service &io_service, const NodeID &self
           new DefaultAgentManagerServiceHandler(agent_manager_)),
       agent_manager_service_(io_service, *agent_manager_service_handler_),
       client_call_manager_(io_service),
+      worker_rpc_pool_(client_call_manager_),
+      local_object_manager_(RayConfig::instance().free_objects_batch_size(),
+                            RayConfig::instance().free_objects_period_milliseconds(),
+                            worker_pool_, gcs_client_->Objects(), worker_rpc_pool_,
+                            [this](const std::vector<ObjectID> &object_ids) {
+                              object_manager_.FreeObjects(object_ids,
+                                                          /*local_only=*/false);
+                            }),
       new_scheduler_enabled_(RayConfig::instance().new_scheduler_enabled()),
       report_worker_backlog_(RayConfig::instance().report_worker_backlog()) {
   RAY_LOG(INFO) << "Initializing NodeManager with ID " << self_node_id_;
@@ -302,7 +309,6 @@ ray::Status NodeManager::RegisterGcs() {
   // Start sending heartbeats to the GCS.
   last_heartbeat_at_ms_ = current_time_ms();
   last_debug_dump_at_ms_ = current_time_ms();
-  last_free_objects_at_ms_ = current_time_ms();
   Heartbeat();
   // Start the timer that gets object manager profiling information and sends it
   // to the GCS.
@@ -495,10 +501,7 @@ void NodeManager::Heartbeat() {
   }
 
   // Evict all copies of freed objects from the cluster.
-  if (free_objects_period_ > 0 &&
-      static_cast<int64_t>(now_ms - last_free_objects_at_ms_) > free_objects_period_) {
-    FlushObjectsToFree();
-  }
+  local_object_manager_.FlushFreeObjectsIfNeeded(now_ms);
 
   // Reset the timer.
   heartbeat_timer_.expires_from_now(heartbeat_period_);
@@ -529,95 +532,14 @@ void NodeManager::DoLocalGC() {
 void NodeManager::HandleRequestObjectSpillage(
     const rpc::RequestObjectSpillageRequest &request,
     rpc::RequestObjectSpillageReply *reply, rpc::SendReplyCallback send_reply_callback) {
-  SpillObjects({ObjectID::FromBinary(request.object_id())},
-               [reply, send_reply_callback](const ray::Status &status) {
-                 if (status.ok()) {
-                   reply->set_success(true);
-                 }
-                 send_reply_callback(Status::OK(), nullptr, nullptr);
-               });
-}
-
-void NodeManager::SpillObjects(const std::vector<ObjectID> &objects_ids,
-                               std::function<void(const ray::Status &)> callback) {
-  for (const auto &id : objects_ids) {
-    // We should not spill an object that we are not the primary copy for.
-    // TODO(swang): We should really return an error here but right now there
-    // is a race condition where the raylet receives the owner's request to
-    // spill an object before it receives the message to pin the objects from
-    // the local worker.
-    if (pinned_objects_.count(id) == 0) {
-      RAY_LOG(WARNING) << "Requested spill for object that has not yet been marked as "
-                          "the primary copy.";
-    }
-  }
-  worker_pool_.PopIOWorker([this, objects_ids,
-                            callback](std::shared_ptr<WorkerInterface> io_worker) {
-    rpc::SpillObjectsRequest request;
-    for (const auto &object_id : objects_ids) {
-      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, objects_ids, callback, io_worker](
-                     const ray::Status &status, const rpc::SpillObjectsReply &r) {
-          worker_pool_.PushIOWorker(io_worker);
-          if (!status.ok()) {
-            RAY_LOG(ERROR) << "Failed to send object spilling request: "
-                           << status.ToString();
-            if (callback) {
-              callback(status);
-            }
-          } else {
-            for (size_t i = 0; i < objects_ids.size(); ++i) {
-              const ObjectID &object_id = objects_ids[i];
-              const std::string &object_url = r.spilled_objects_url(i);
-              RAY_LOG(DEBUG) << "Object " << object_id << " spilled at " << object_url;
-              // Write to object directory. Wait for the write to finish before
-              // releasing the object to make sure that the spilled object can
-              // be retrieved by other raylets.
-              RAY_CHECK_OK(gcs_client_->Objects().AsyncAddSpilledUrl(
-                  object_id, object_url, [this, object_id, callback](Status status) {
-                    RAY_CHECK_OK(status);
-                    // 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);
-                    if (callback) {
-                      callback(status);
-                    }
-                  }));
-            }
-          }
-        });
-  });
-}
-
-void NodeManager::AsyncRestoreSpilledObject(
-    const ObjectID &object_id, const std::string &object_url,
-    std::function<void(const ray::Status &)> callback) {
-  RAY_LOG(DEBUG) << "Restoring spilled object " << object_id << " from URL "
-                 << object_url;
-  worker_pool_.PopIOWorker([this, object_url,
-                            callback](std::shared_ptr<WorkerInterface> 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) {
-          worker_pool_.PushIOWorker(io_worker);
-          if (!status.ok()) {
-            RAY_LOG(ERROR) << "Failed to send restore spilled object request: "
-                           << status.ToString();
-          }
-          if (callback) {
-            callback(status);
-          }
-        });
-  });
+  local_object_manager_.SpillObjects(
+      {ObjectID::FromBinary(request.object_id())},
+      [reply, send_reply_callback](const ray::Status &status) {
+        if (status.ok()) {
+          reply->set_success(true);
+        }
+        send_reply_callback(Status::OK(), nullptr, nullptr);
+      });
 }
 
 // TODO(edoakes): this function is problematic because it both sends warnings spuriously
@@ -3065,28 +2987,16 @@ std::string compact_tag_string(const opencensus::stats::ViewDescriptor &view,
 void NodeManager::HandlePinObjectIDs(const rpc::PinObjectIDsRequest &request,
                                      rpc::PinObjectIDsReply *reply,
                                      rpc::SendReplyCallback send_reply_callback) {
-  WorkerID worker_id = WorkerID::FromBinary(request.owner_address().worker_id());
-  auto it = worker_rpc_clients_.find(worker_id);
-  if (it == worker_rpc_clients_.end()) {
-    auto client = std::unique_ptr<rpc::CoreWorkerClient>(
-        new rpc::CoreWorkerClient(request.owner_address(), client_call_manager_));
-    it = worker_rpc_clients_
-             .emplace(worker_id,
-                      std::make_pair<std::unique_ptr<rpc::CoreWorkerClient>, size_t>(
-                          std::move(client), 0))
-             .first;
+  std::vector<ObjectID> object_ids;
+  object_ids.reserve(request.object_ids_size());
+  for (const auto &object_id_binary : request.object_ids()) {
+    object_ids.push_back(ObjectID::FromBinary(object_id_binary));
   }
-
   if (object_pinning_enabled_) {
     // Pin the objects in plasma by getting them and holding a reference to
     // the returned buffer.
     // NOTE: the caller must ensure that the objects already exist in plasma before
     // sending a PinObjectIDs request.
-    std::vector<ObjectID> object_ids;
-    object_ids.reserve(request.object_ids_size());
-    for (const auto &object_id_binary : request.object_ids()) {
-      object_ids.push_back(ObjectID::FromBinary(object_id_binary));
-    }
     std::vector<plasma::ObjectBuffer> plasma_results;
     // TODO(swang): This `Get` has a timeout of 0, so the plasma store will not
     // block when serving the request. However, if the plasma store is under
@@ -3100,77 +3010,23 @@ void NodeManager::HandlePinObjectIDs(const rpc::PinObjectIDsRequest &request,
       return;
     }
 
-    // Pin the requested objects until the owner notifies us that the objects can be
-    // unpinned by responding to the WaitForObjectEviction message.
-    // TODO(edoakes): we should be batching these requests instead of sending one per
-    // pinned object.
+    std::vector<std::unique_ptr<RayObject>> objects;
     for (int64_t i = 0; i < request.object_ids().size(); i++) {
-      ObjectID object_id = ObjectID::FromBinary(request.object_ids(i));
-
       if (plasma_results[i].data == nullptr) {
-        RAY_LOG(ERROR) << "Plasma object " << object_id
-                       << " was evicted before the raylet could pin it.";
-        continue;
+        objects.push_back(nullptr);
+      } else {
+        objects.emplace_back(std::unique_ptr<RayObject>(new RayObject(
+            std::make_shared<PlasmaBuffer>(plasma_results[i].data),
+            std::make_shared<PlasmaBuffer>(plasma_results[i].metadata), {})));
       }
-
-      RAY_LOG(DEBUG) << "Pinning object " << object_id;
-      RAY_CHECK(
-          pinned_objects_
-              .emplace(
-                  object_id,
-                  std::unique_ptr<RayObject>(new RayObject(
-                      std::make_shared<PlasmaBuffer>(plasma_results[i].data),
-                      std::make_shared<PlasmaBuffer>(plasma_results[i].metadata), {})))
-              .second);
     }
+    local_object_manager_.PinObjects(object_ids, std::move(objects));
   }
-
-  for (const auto &object_id_binary : request.object_ids()) {
-    ObjectID object_id = ObjectID::FromBinary(object_id_binary);
-    // Send a long-running RPC request to the owner for each object. When we get a
-    // response or the RPC fails (due to the owner crashing), unpin the object.
-    rpc::WaitForObjectEvictionRequest wait_request;
-    wait_request.set_object_id(object_id_binary);
-    wait_request.set_intended_worker_id(request.owner_address().worker_id());
-    worker_rpc_clients_[worker_id].second++;
-    it->second.first->WaitForObjectEviction(
-        wait_request, [this, worker_id, object_id](
-                          Status status, const rpc::WaitForObjectEvictionReply &reply) {
-          if (!status.ok()) {
-            RAY_LOG(WARNING) << "Worker " << worker_id << " failed. Unpinning object "
-                             << object_id;
-          }
-          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_ >= 0) {
-            objects_to_free_.push_back(object_id);
-          }
-          if (objects_to_free_.size() ==
-                  RayConfig::instance().free_objects_batch_size() ||
-              free_objects_period_ == 0) {
-            FlushObjectsToFree();
-          }
-
-          // Remove the cached worker client if there are no more pending requests.
-          if (--worker_rpc_clients_[worker_id].second == 0) {
-            worker_rpc_clients_.erase(worker_id);
-          }
-        });
-  }
+  // Wait for the object to be freed by the owner, which keeps the ref count.
+  local_object_manager_.WaitForObjectFree(request.owner_address(), object_ids);
   send_reply_callback(Status::OK(), nullptr, nullptr);
 }
 
-void NodeManager::FlushObjectsToFree() {
-  if (!objects_to_free_.empty()) {
-    RAY_LOG(DEBUG) << "Freeing " << objects_to_free_.size() << " out-of-scope objects";
-    object_manager_.FreeObjects(objects_to_free_, /*local_only=*/false);
-    objects_to_free_.clear();
-  }
-  last_free_objects_at_ms_ = current_time_ms();
-}
-
 void NodeManager::HandleGetNodeStats(const rpc::GetNodeStatsRequest &node_stats_request,
                                      rpc::GetNodeStatsReply *reply,
                                      rpc::SendReplyCallback send_reply_callback) {

src/ray/raylet/node_manager.h

@@ -28,6 +28,7 @@
 #include "ray/object_manager/object_manager.h"
 #include "ray/raylet/actor_registration.h"
 #include "ray/raylet/agent_manager.h"
+#include "ray/raylet/local_object_manager.h"
 #include "ray/raylet/scheduling/scheduling_ids.h"
 #include "ray/raylet/scheduling/cluster_resource_scheduler.h"
 #include "ray/raylet/scheduling/cluster_task_manager.h"
@@ -36,6 +37,7 @@
 #include "ray/raylet/reconstruction_policy.h"
 #include "ray/raylet/task_dependency_manager.h"
 #include "ray/raylet/worker_pool.h"
+#include "ray/rpc/worker/core_worker_client_pool.h"
 #include "ray/util/ordered_set.h"
 #include "ray/common/bundle_spec.h"
 // clang-format on
@@ -85,9 +87,6 @@ struct NodeManagerConfig {
   uint64_t heartbeat_period_ms;
   /// The time between debug dumps in milliseconds, or -1 to disable.
   uint64_t debug_dump_period_ms;
-  /// The time between attempts to eagerly evict objects from plasma in
-  /// milliseconds, or -1 to disable.
-  int64_t free_objects_period_ms;
   /// Whether to enable fair queueing between task classes in raylet.
   bool fair_queueing_enabled;
   /// Whether to enable pinning for plasma objects.
@@ -171,13 +170,7 @@ class NodeManager : public rpc::NodeManagerServiceHandler {
   /// Get the port of the node manager rpc server.
   int GetServerPort() const { return node_manager_server_.GetPort(); }
 
-  /// Restore a spilled object from external storage back into local memory.
-  /// \param object_id The ID of the object to restore.
-  /// \param object_url The URL in external storage from which the object can be restored.
-  /// \param callback A callback to call when the restoration is done. Status
-  /// will contain the error during restoration, if any.
-  void AsyncRestoreSpilledObject(const ObjectID &object_id, const std::string &object_url,
-                                 std::function<void(const ray::Status &)> callback);
+  LocalObjectManager &GetLocalObjectManager() { return local_object_manager_; }
 
  private:
   /// Methods for handling clients.
@@ -656,11 +649,6 @@ class NodeManager : public rpc::NodeManagerServiceHandler {
   /// Trigger local GC on each worker of this raylet.
   void DoLocalGC();
 
-  /// Spill objects to external storage.
-  /// \param objects_ids_to_spill The objects to be spilled.
-  void SpillObjects(const std::vector<ObjectID> &objects_ids_to_spill,
-                    std::function<void(const ray::Status &)> callback = nullptr);
-
   /// Push an error to the driver if this node is full of actors and so we are
   /// unable to schedule new tasks or actors at all.
   void WarnResourceDeadlock();
@@ -702,8 +690,6 @@ class NodeManager : public rpc::NodeManagerServiceHandler {
   std::chrono::milliseconds heartbeat_period_;
   /// The period between debug state dumps.
   int64_t debug_dump_period_;
-  /// The period between attempts to eagerly evict objects from plasma.
-  int64_t free_objects_period_;
   /// Whether to enable fair queueing between task classes in raylet.
   bool fair_queueing_enabled_;
   /// Whether to enable pinning for plasma objects.
@@ -727,9 +713,6 @@ class NodeManager : public rpc::NodeManagerServiceHandler {
   SchedulingResources last_heartbeat_resources_;
   /// The time that the last debug string was logged to the console.
   uint64_t last_debug_dump_at_ms_;
-  /// The time that we last sent a FreeObjects request to other nodes for
-  /// objects that have gone out of scope in the application.
-  uint64_t last_free_objects_at_ms_;
   /// The number of heartbeats that we should wait before sending the
   /// next load report.
   uint8_t num_heartbeats_before_load_report_;
@@ -772,6 +755,13 @@ class NodeManager : public rpc::NodeManagerServiceHandler {
   /// as well as all `CoreWorkerClient`s.
   rpc::ClientCallManager client_call_manager_;
 
+  /// Pool of RPC client connections to core workers.
+  rpc::CoreWorkerClientPool worker_rpc_pool_;
+
+  /// Manages all local objects that are pinned (primary
+  /// copies), freed, and/or spilled.
+  LocalObjectManager local_object_manager_;
+
   /// Map from node ids to clients of the remote node managers.
   std::unordered_map<NodeID, std::unique_ptr<rpc::NodeManagerClient>>
       remote_node_manager_clients_;
@@ -804,13 +794,6 @@ class NodeManager : public rpc::NodeManagerServiceHandler {
   std::shared_ptr<ClusterResourceScheduler> new_resource_scheduler_;
   std::shared_ptr<ClusterTaskManager> cluster_task_manager_;
 
-  /// Cache of gRPC clients to workers (not necessarily running on this node).
-  /// Also includes the number of inflight requests to each worker - when this
-  /// reaches zero, the client will be deleted and a new one will need to be created
-  /// for any subsequent requests.
-  absl::flat_hash_map<WorkerID, std::pair<std::unique_ptr<rpc::CoreWorkerClient>, size_t>>
-      worker_rpc_clients_;
-
   absl::flat_hash_map<ObjectID, std::unique_ptr<RayObject>> pinned_objects_;
 
   // TODO(swang): Evict entries from these caches.
@@ -826,12 +809,6 @@ class NodeManager : public rpc::NodeManagerServiceHandler {
   absl::flat_hash_map<ObjectID, absl::flat_hash_set<std::shared_ptr<WorkerInterface>>>
       async_plasma_objects_notification_ GUARDED_BY(plasma_object_notification_lock_);
 
-  /// Objects that are out of scope in the application and that should be freed
-  /// from plasma. The cache is flushed when it reaches the config's
-  /// 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<ObjectID> objects_to_free_;
-
   /// This map represents the commit state of 2PC protocol for atomic placement group
   /// creation.
   absl::flat_hash_map<BundleID, std::shared_ptr<BundleState>, pair_hash>

src/ray/raylet/raylet.cc

@@ -69,12 +69,13 @@ Raylet::Raylet(boost::asio::io_service &main_service, const std::string &socket_
                                                                     gcs_client_))
               : std::dynamic_pointer_cast<ObjectDirectoryInterface>(
                     std::make_shared<ObjectDirectory>(main_service, gcs_client_))),
-      object_manager_(
-          main_service, self_node_id_, object_manager_config, object_directory_,
-          [this](const ObjectID &object_id, const std::string &spilled_url,
-                 std::function<void(const ray::Status &)> callback) {
-            node_manager_.AsyncRestoreSpilledObject(object_id, spilled_url, callback);
-          }),
+      object_manager_(main_service, self_node_id_, object_manager_config,
+                      object_directory_,
+                      [this](const ObjectID &object_id, const std::string &spilled_url,
+                             std::function<void(const ray::Status &)> callback) {
+                        node_manager_.GetLocalObjectManager().AsyncRestoreSpilledObject(
+                            object_id, spilled_url, callback);
+                      }),
       node_manager_(main_service, self_node_id_, node_manager_config, object_manager_,
                     gcs_client_, object_directory_),
       socket_name_(socket_name),

src/ray/raylet/scheduling/cluster_task_manager_test.cc

@@ -27,6 +27,7 @@
 #include "ray/common/test_util.h"
 #include "ray/raylet/scheduling/cluster_resource_scheduler.h"
 #include "ray/raylet/scheduling/scheduling_ids.h"
+#include "ray/raylet/test/util.h"
 
 #ifdef UNORDERED_VS_ABSL_MAPS_EVALUATION
 #include <chrono>
@@ -56,195 +57,6 @@ class MockWorkerPool : public WorkerPoolInterface {
   std::list<std::shared_ptr<WorkerInterface>> workers;
 };
 
-class MockWorker : public WorkerInterface {
- public:
-  MockWorker(WorkerID worker_id, int port) : worker_id_(worker_id), port_(port) {}
-
-  WorkerID WorkerId() const { return worker_id_; }
-
-  rpc::WorkerType GetWorkerType() const { return rpc::WorkerType::WORKER; }
-
-  int Port() const { return port_; }
-
-  void SetOwnerAddress(const rpc::Address &address) { address_ = address; }
-
-  void AssignTaskId(const TaskID &task_id) {}
-
-  void AssignJobId(const JobID &job_id) {}
-
-  void SetAssignedTask(Task &assigned_task) {}
-
-  const std::string IpAddress() const { return address_.ip_address(); }
-
-  void SetAllocatedInstances(
-      std::shared_ptr<TaskResourceInstances> &allocated_instances) {
-    allocated_instances_ = allocated_instances;
-  }
-
-  void SetLifetimeAllocatedInstances(
-      std::shared_ptr<TaskResourceInstances> &allocated_instances) {
-    lifetime_allocated_instances_ = allocated_instances;
-  }
-
-  std::shared_ptr<TaskResourceInstances> GetAllocatedInstances() {
-    return allocated_instances_;
-  }
-  std::shared_ptr<TaskResourceInstances> GetLifetimeAllocatedInstances() {
-    return lifetime_allocated_instances_;
-  }
-
-  void MarkDead() { RAY_CHECK(false) << "Method unused"; }
-  bool IsDead() const {
-    RAY_CHECK(false) << "Method unused";
-    return false;
-  }
-  void MarkBlocked() { RAY_CHECK(false) << "Method unused"; }
-  void MarkUnblocked() { RAY_CHECK(false) << "Method unused"; }
-  bool IsBlocked() const {
-    RAY_CHECK(false) << "Method unused";
-    return false;
-  }
-
-  Process GetProcess() const {
-    RAY_CHECK(false) << "Method unused";
-    return Process::CreateNewDummy();
-  }
-  void SetProcess(Process proc) { RAY_CHECK(false) << "Method unused"; }
-  Language GetLanguage() const {
-    RAY_CHECK(false) << "Method unused";
-    return Language::PYTHON;
-  }
-
-  void Connect(int port) { RAY_CHECK(false) << "Method unused"; }
-
-  int AssignedPort() const {
-    RAY_CHECK(false) << "Method unused";
-    return -1;
-  }
-  void SetAssignedPort(int port) { RAY_CHECK(false) << "Method unused"; }
-  const TaskID &GetAssignedTaskId() const {
-    RAY_CHECK(false) << "Method unused";
-    return TaskID::Nil();
-  }
-  bool AddBlockedTaskId(const TaskID &task_id) {
-    RAY_CHECK(false) << "Method unused";
-    return false;
-  }
-  bool RemoveBlockedTaskId(const TaskID &task_id) {
-    RAY_CHECK(false) << "Method unused";
-    return false;
-  }
-  const std::unordered_set<TaskID> &GetBlockedTaskIds() const {
-    RAY_CHECK(false) << "Method unused";
-    auto *t = new std::unordered_set<TaskID>();
-    return *t;
-  }
-  const JobID &GetAssignedJobId() const {
-    RAY_CHECK(false) << "Method unused";
-    return JobID::Nil();
-  }
-  void AssignActorId(const ActorID &actor_id) { RAY_CHECK(false) << "Method unused"; }
-  const ActorID &GetActorId() const {
-    RAY_CHECK(false) << "Method unused";
-    return ActorID::Nil();
-  }
-  void MarkDetachedActor() { RAY_CHECK(false) << "Method unused"; }
-  bool IsDetachedActor() const {
-    RAY_CHECK(false) << "Method unused";
-    return false;
-  }
-  const std::shared_ptr<ClientConnection> Connection() const {
-    RAY_CHECK(false) << "Method unused";
-    return nullptr;
-  }
-  const rpc::Address &GetOwnerAddress() const {
-    RAY_CHECK(false) << "Method unused";
-    return address_;
-  }
-
-  const ResourceIdSet &GetLifetimeResourceIds() const {
-    RAY_CHECK(false) << "Method unused";
-    auto *t = new ResourceIdSet();
-    return *t;
-  }
-  void SetLifetimeResourceIds(ResourceIdSet &resource_ids) {
-    RAY_CHECK(false) << "Method unused";
-  }
-  void ResetLifetimeResourceIds() { RAY_CHECK(false) << "Method unused"; }
-
-  const ResourceIdSet &GetTaskResourceIds() const {
-    RAY_CHECK(false) << "Method unused";
-    auto *t = new ResourceIdSet();
-    return *t;
-  }
-  void SetTaskResourceIds(ResourceIdSet &resource_ids) {
-    RAY_CHECK(false) << "Method unused";
-  }
-  void ResetTaskResourceIds() { RAY_CHECK(false) << "Method unused"; }
-  ResourceIdSet ReleaseTaskCpuResources() {
-    RAY_CHECK(false) << "Method unused";
-    auto *t = new ResourceIdSet();
-    return *t;
-  }
-  void AcquireTaskCpuResources(const ResourceIdSet &cpu_resources) {
-    RAY_CHECK(false) << "Method unused";
-  }
-
-  Status AssignTask(const Task &task, const ResourceIdSet &resource_id_set) {
-    RAY_CHECK(false) << "Method unused";
-    Status s;
-    return s;
-  }
-  void DirectActorCallArgWaitComplete(int64_t tag) {
-    RAY_CHECK(false) << "Method unused";
-  }
-
-  void ClearAllocatedInstances() { allocated_instances_ = nullptr; }
-
-  void ClearLifetimeAllocatedInstances() { RAY_CHECK(false) << "Method unused"; }
-
-  void SetBorrowedCPUInstances(std::vector<double> &cpu_instances) {
-    borrowed_cpu_instances_ = cpu_instances;
-  }
-
-  const PlacementGroupID &GetPlacementGroupId() const {
-    RAY_CHECK(false) << "Method unused";
-    return PlacementGroupID::Nil();
-  }
-
-  void SetPlacementGroupId(const PlacementGroupID &placement_group_id) {
-    RAY_CHECK(false) << "Method unused";
-  }
-
-  std::vector<double> &GetBorrowedCPUInstances() { return borrowed_cpu_instances_; }
-
-  void ClearBorrowedCPUInstances() { RAY_CHECK(false) << "Method unused"; }
-
-  Task &GetAssignedTask() {
-    RAY_CHECK(false) << "Method unused";
-    auto *t = new Task();
-    return *t;
-  }
-
-  bool IsRegistered() {
-    RAY_CHECK(false) << "Method unused";
-    return false;
-  }
-
-  rpc::CoreWorkerClient *rpc_client() {
-    RAY_CHECK(false) << "Method unused";
-    return nullptr;
-  }
-
- private:
-  WorkerID worker_id_;
-  int port_;
-  rpc::Address address_;
-  std::shared_ptr<TaskResourceInstances> allocated_instances_;
-  std::shared_ptr<TaskResourceInstances> lifetime_allocated_instances_;
-  std::vector<double> borrowed_cpu_instances_;
-};
-
 std::shared_ptr<ClusterResourceScheduler> CreateSingleNodeScheduler(
     const std::string &id) {
   std::unordered_map<std::string, double> local_node_resources;

src/ray/raylet/test/local_object_manager_test.cc

@@ -0,0 +1,259 @@
+// Copyright 2017 The Ray Authors.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//  http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include "ray/raylet/local_object_manager.h"
+
+#include "gmock/gmock.h"
+#include "gtest/gtest.h"
+#include "ray/common/id.h"
+#include "ray/gcs/accessor.h"
+#include "ray/raylet/test/util.h"
+#include "ray/raylet/worker_pool.h"
+#include "ray/rpc/grpc_client.h"
+#include "ray/rpc/worker/core_worker_client.h"
+#include "ray/rpc/worker/core_worker_client_pool.h"
+#include "src/ray/protobuf/core_worker.grpc.pb.h"
+#include "src/ray/protobuf/core_worker.pb.h"
+
+namespace ray {
+
+namespace raylet {
+
+using ::testing::_;
+
+class MockWorkerClient : public rpc::CoreWorkerClientInterface {
+ public:
+  void WaitForObjectEviction(
+      const rpc::WaitForObjectEvictionRequest &request,
+      const rpc::ClientCallback<rpc::WaitForObjectEvictionReply> &callback) override {
+    callbacks.push_back(callback);
+  }
+
+  bool ReplyObjectEviction(Status status = Status::OK()) {
+    if (callbacks.size() == 0) {
+      return false;
+    }
+    auto callback = callbacks.front();
+    auto reply = rpc::WaitForObjectEvictionReply();
+    callback(status, reply);
+    callbacks.pop_front();
+    return true;
+  }
+
+  std::list<rpc::ClientCallback<rpc::WaitForObjectEvictionReply>> callbacks;
+};
+
+class MockIOWorkerClient : public rpc::CoreWorkerClientInterface {
+ public:
+  void SpillObjects(
+      const rpc::SpillObjectsRequest &request,
+      const rpc::ClientCallback<rpc::SpillObjectsReply> &callback) override {
+    callbacks.push_back(callback);
+  }
+
+  bool ReplySpillObjects(std::vector<std::string> urls, Status status = Status::OK()) {
+    if (callbacks.size() == 0) {
+      return false;
+    }
+    auto callback = callbacks.front();
+    auto reply = rpc::SpillObjectsReply();
+    for (const auto &url : urls) {
+      reply.add_spilled_objects_url(url);
+    }
+    callback(status, reply);
+    callbacks.pop_front();
+    return true;
+  }
+
+  void RestoreSpilledObjects(
+      const rpc::RestoreSpilledObjectsRequest &request,
+      const rpc::ClientCallback<rpc::RestoreSpilledObjectsReply> &callback) override {
+    rpc::RestoreSpilledObjectsReply reply;
+    callback(Status(), reply);
+  }
+
+  std::list<rpc::ClientCallback<rpc::SpillObjectsReply>> callbacks;
+};
+
+class MockIOWorker : public MockWorker {
+ public:
+  MockIOWorker(WorkerID worker_id, int port,
+               std::shared_ptr<rpc::CoreWorkerClientInterface> io_worker)
+      : MockWorker(worker_id, port), io_worker(io_worker) {}
+
+  rpc::CoreWorkerClientInterface *rpc_client() { return io_worker.get(); }
+
+  std::shared_ptr<rpc::CoreWorkerClientInterface> io_worker;
+};
+
+class MockIOWorkerPool : public IOWorkerPoolInterface {
+ public:
+  MOCK_METHOD1(PushIOWorker, void(const std::shared_ptr<WorkerInterface> &worker));
+
+  void PopIOWorker(
+      std::function<void(std::shared_ptr<WorkerInterface>)> callback) override {
+    callback(io_worker);
+  }
+
+  std::shared_ptr<MockIOWorkerClient> io_worker_client =
+      std::make_shared<MockIOWorkerClient>();
+  std::shared_ptr<WorkerInterface> io_worker =
+      std::make_shared<MockIOWorker>(WorkerID::FromRandom(), 1234, io_worker_client);
+};
+
+class MockObjectInfoAccessor : public gcs::ObjectInfoAccessor {
+ public:
+  MOCK_METHOD2(
+      AsyncGetLocations,
+      Status(const ObjectID &object_id,
+             const gcs::OptionalItemCallback<rpc::ObjectLocationInfo> &callback));
+
+  MOCK_METHOD1(AsyncGetAll,
+               Status(const gcs::MultiItemCallback<rpc::ObjectLocationInfo> &callback));
+
+  MOCK_METHOD3(AsyncAddLocation, Status(const ObjectID &object_id, const NodeID &node_id,
+                                        const gcs::StatusCallback &callback));
+
+  Status AsyncAddSpilledUrl(const ObjectID &object_id, const std::string &spilled_url,
+                            const gcs::StatusCallback &callback) {
+    object_urls[object_id] = spilled_url;
+    callback(Status());
+    return Status();
+  }
+
+  MOCK_METHOD3(AsyncRemoveLocation,
+               Status(const ObjectID &object_id, const NodeID &node_id,
+                      const gcs::StatusCallback &callback));
+
+  MOCK_METHOD3(AsyncSubscribeToLocations,
+               Status(const ObjectID &object_id,
+                      const gcs::SubscribeCallback<
+                          ObjectID, std::vector<rpc::ObjectLocationChange>> &subscribe,
+                      const gcs::StatusCallback &done));
+
+  MOCK_METHOD1(AsyncUnsubscribeToLocations, Status(const ObjectID &object_id));
+
+  MOCK_METHOD1(AsyncResubscribe, void(bool is_pubsub_server_restarted));
+
+  MOCK_METHOD1(IsObjectUnsubscribed, bool(const ObjectID &object_id));
+
+  std::unordered_map<ObjectID, std::string> object_urls;
+};
+
+class LocalObjectManagerTest : public ::testing::Test {
+ public:
+  LocalObjectManagerTest()
+      : owner_client(std::make_shared<MockWorkerClient>()),
+        client_pool([&](const rpc::Address &addr) { return owner_client; }),
+        manager(free_objects_batch_size,
+                /*free_objects_period_ms=*/1000, worker_pool, object_table, client_pool,
+                [&](const std::vector<ObjectID> &object_ids) {
+                  for (const auto &object_id : object_ids) {
+                    freed.insert(object_id);
+                  }
+                }) {}
+
+  size_t free_objects_batch_size = 3;
+  std::shared_ptr<MockWorkerClient> owner_client;
+  rpc::CoreWorkerClientPool client_pool;
+  MockIOWorkerPool worker_pool;
+  MockObjectInfoAccessor object_table;
+  LocalObjectManager manager;
+
+  std::unordered_set<ObjectID> freed;
+};
+
+TEST_F(LocalObjectManagerTest, TestPin) {
+  rpc::Address owner_address;
+  owner_address.set_worker_id(WorkerID::FromRandom().Binary());
+
+  std::vector<ObjectID> object_ids;
+  std::vector<std::unique_ptr<RayObject>> objects;
+
+  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<int>(rpc::ErrorType::OBJECT_IN_PLASMA));
+    auto metadata = const_cast<uint8_t *>(reinterpret_cast<const uint8_t *>(meta.data()));
+    auto meta_buffer = std::make_shared<LocalMemoryBuffer>(metadata, meta.size());
+    std::unique_ptr<RayObject> object(
+        new RayObject(nullptr, meta_buffer, std::vector<ObjectID>()));
+    objects.push_back(std::move(object));
+  }
+  manager.PinObjects(object_ids, std::move(objects));
+  manager.WaitForObjectFree(owner_address, object_ids);
+
+  for (size_t i = 0; i < free_objects_batch_size; i++) {
+    ASSERT_TRUE(freed.empty());
+    ASSERT_TRUE(owner_client->ReplyObjectEviction());
+  }
+  std::unordered_set<ObjectID> expected(object_ids.begin(), object_ids.end());
+  ASSERT_EQ(freed, expected);
+}
+
+TEST_F(LocalObjectManagerTest, TestRestoreSpilledObject) {
+  ObjectID object_id = ObjectID::FromRandom();
+  std::string object_url("url");
+  int num_times_fired = 0;
+  EXPECT_CALL(worker_pool, PushIOWorker(_));
+  manager.AsyncRestoreSpilledObject(object_id, object_url, [&](const Status &status) {
+    ASSERT_TRUE(status.ok());
+    num_times_fired++;
+  });
+  ASSERT_EQ(num_times_fired, 1);
+}
+
+TEST_F(LocalObjectManagerTest, TestExplicitSpill) {
+  std::vector<ObjectID> object_ids;
+  std::vector<std::unique_ptr<RayObject>> objects;
+
+  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<int>(rpc::ErrorType::OBJECT_IN_PLASMA));
+    auto metadata = const_cast<uint8_t *>(reinterpret_cast<const uint8_t *>(meta.data()));
+    auto meta_buffer = std::make_shared<LocalMemoryBuffer>(metadata, meta.size());
+    std::unique_ptr<RayObject> object(
+        new RayObject(nullptr, meta_buffer, std::vector<ObjectID>()));
+    objects.push_back(std::move(object));
+  }
+  manager.PinObjects(object_ids, std::move(objects));
+
+  int num_times_fired = 0;
+  manager.SpillObjects(object_ids, [&](const Status &status) mutable {
+    ASSERT_TRUE(status.ok());
+    num_times_fired++;
+  });
+  ASSERT_EQ(num_times_fired, 0);
+
+  EXPECT_CALL(worker_pool, PushIOWorker(_));
+  std::vector<std::string> urls;
+  for (size_t i = 0; i < object_ids.size(); i++) {
+    urls.push_back("url" + std::to_string(i));
+  }
+  ASSERT_TRUE(worker_pool.io_worker_client->ReplySpillObjects(urls));
+  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]);
+  }
+}
+
+}  // namespace raylet
+
+}  // namespace ray
+
+int main(int argc, char **argv) {
+  ::testing::InitGoogleTest(&argc, argv);
+  return RUN_ALL_TESTS();
+}

src/ray/raylet/test/util.h

@@ -0,0 +1,212 @@
+// Copyright 2017 The Ray Authors.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//  http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include "ray/raylet/worker.h"
+
+namespace ray {
+
+namespace raylet {
+
+class MockWorker : public WorkerInterface {
+ public:
+  MockWorker(WorkerID worker_id, int port) : worker_id_(worker_id), port_(port) {}
+
+  WorkerID WorkerId() const { return worker_id_; }
+
+  rpc::WorkerType GetWorkerType() const { return rpc::WorkerType::WORKER; }
+
+  int Port() const { return port_; }
+
+  void SetOwnerAddress(const rpc::Address &address) { address_ = address; }
+
+  void AssignTaskId(const TaskID &task_id) {}
+
+  void AssignJobId(const JobID &job_id) {}
+
+  void SetAssignedTask(Task &assigned_task) {}
+
+  const std::string IpAddress() const { return address_.ip_address(); }
+
+  void SetAllocatedInstances(
+      std::shared_ptr<TaskResourceInstances> &allocated_instances) {
+    allocated_instances_ = allocated_instances;
+  }
+
+  void SetLifetimeAllocatedInstances(
+      std::shared_ptr<TaskResourceInstances> &allocated_instances) {
+    lifetime_allocated_instances_ = allocated_instances;
+  }
+
+  std::shared_ptr<TaskResourceInstances> GetAllocatedInstances() {
+    return allocated_instances_;
+  }
+  std::shared_ptr<TaskResourceInstances> GetLifetimeAllocatedInstances() {
+    return lifetime_allocated_instances_;
+  }
+
+  void MarkDead() { RAY_CHECK(false) << "Method unused"; }
+  bool IsDead() const {
+    RAY_CHECK(false) << "Method unused";
+    return false;
+  }
+  void MarkBlocked() { RAY_CHECK(false) << "Method unused"; }
+  void MarkUnblocked() { RAY_CHECK(false) << "Method unused"; }
+  bool IsBlocked() const {
+    RAY_CHECK(false) << "Method unused";
+    return false;
+  }
+
+  Process GetProcess() const {
+    RAY_CHECK(false) << "Method unused";
+    return Process::CreateNewDummy();
+  }
+  void SetProcess(Process proc) { RAY_CHECK(false) << "Method unused"; }
+  Language GetLanguage() const {
+    RAY_CHECK(false) << "Method unused";
+    return Language::PYTHON;
+  }
+
+  void Connect(int port) { RAY_CHECK(false) << "Method unused"; }
+
+  int AssignedPort() const {
+    RAY_CHECK(false) << "Method unused";
+    return -1;
+  }
+  void SetAssignedPort(int port) { RAY_CHECK(false) << "Method unused"; }
+  const TaskID &GetAssignedTaskId() const {
+    RAY_CHECK(false) << "Method unused";
+    return TaskID::Nil();
+  }
+  bool AddBlockedTaskId(const TaskID &task_id) {
+    RAY_CHECK(false) << "Method unused";
+    return false;
+  }
+  bool RemoveBlockedTaskId(const TaskID &task_id) {
+    RAY_CHECK(false) << "Method unused";
+    return false;
+  }
+  const std::unordered_set<TaskID> &GetBlockedTaskIds() const {
+    RAY_CHECK(false) << "Method unused";
+    auto *t = new std::unordered_set<TaskID>();
+    return *t;
+  }
+  const JobID &GetAssignedJobId() const {
+    RAY_CHECK(false) << "Method unused";
+    return JobID::Nil();
+  }
+  void AssignActorId(const ActorID &actor_id) { RAY_CHECK(false) << "Method unused"; }
+  const ActorID &GetActorId() const {
+    RAY_CHECK(false) << "Method unused";
+    return ActorID::Nil();
+  }
+  void MarkDetachedActor() { RAY_CHECK(false) << "Method unused"; }
+  bool IsDetachedActor() const {
+    RAY_CHECK(false) << "Method unused";
+    return false;
+  }
+  const std::shared_ptr<ClientConnection> Connection() const {
+    RAY_CHECK(false) << "Method unused";
+    return nullptr;
+  }
+  const rpc::Address &GetOwnerAddress() const {
+    RAY_CHECK(false) << "Method unused";
+    return address_;
+  }
+
+  const ResourceIdSet &GetLifetimeResourceIds() const {
+    RAY_CHECK(false) << "Method unused";
+    auto *t = new ResourceIdSet();
+    return *t;
+  }
+  void SetLifetimeResourceIds(ResourceIdSet &resource_ids) {
+    RAY_CHECK(false) << "Method unused";
+  }
+  void ResetLifetimeResourceIds() { RAY_CHECK(false) << "Method unused"; }
+
+  const ResourceIdSet &GetTaskResourceIds() const {
+    RAY_CHECK(false) << "Method unused";
+    auto *t = new ResourceIdSet();
+    return *t;
+  }
+  void SetTaskResourceIds(ResourceIdSet &resource_ids) {
+    RAY_CHECK(false) << "Method unused";
+  }
+  void ResetTaskResourceIds() { RAY_CHECK(false) << "Method unused"; }
+  ResourceIdSet ReleaseTaskCpuResources() {
+    RAY_CHECK(false) << "Method unused";
+    auto *t = new ResourceIdSet();
+    return *t;
+  }
+  void AcquireTaskCpuResources(const ResourceIdSet &cpu_resources) {
+    RAY_CHECK(false) << "Method unused";
+  }
+
+  Status AssignTask(const Task &task, const ResourceIdSet &resource_id_set) {
+    RAY_CHECK(false) << "Method unused";
+    Status s;
+    return s;
+  }
+  void DirectActorCallArgWaitComplete(int64_t tag) {
+    RAY_CHECK(false) << "Method unused";
+  }
+
+  void ClearAllocatedInstances() { allocated_instances_ = nullptr; }
+
+  void ClearLifetimeAllocatedInstances() { RAY_CHECK(false) << "Method unused"; }
+
+  void SetBorrowedCPUInstances(std::vector<double> &cpu_instances) {
+    borrowed_cpu_instances_ = cpu_instances;
+  }
+
+  const PlacementGroupID &GetPlacementGroupId() const {
+    RAY_CHECK(false) << "Method unused";
+    return PlacementGroupID::Nil();
+  }
+
+  void SetPlacementGroupId(const PlacementGroupID &placement_group_id) {
+    RAY_CHECK(false) << "Method unused";
+  }
+
+  std::vector<double> &GetBorrowedCPUInstances() { return borrowed_cpu_instances_; }
+
+  void ClearBorrowedCPUInstances() { RAY_CHECK(false) << "Method unused"; }
+
+  Task &GetAssignedTask() {
+    RAY_CHECK(false) << "Method unused";
+    auto *t = new Task();
+    return *t;
+  }
+
+  bool IsRegistered() {
+    RAY_CHECK(false) << "Method unused";
+    return false;
+  }
+
+  rpc::CoreWorkerClientInterface *rpc_client() {
+    RAY_CHECK(false) << "Method unused";
+    return nullptr;
+  }
+
+ private:
+  WorkerID worker_id_;
+  int port_;
+  rpc::Address address_;
+  std::shared_ptr<TaskResourceInstances> allocated_instances_;
+  std::shared_ptr<TaskResourceInstances> lifetime_allocated_instances_;
+  std::vector<double> borrowed_cpu_instances_;
+};
+
+}  // namespace raylet
+
+}  // namespace ray

src/ray/raylet/worker.h

@@ -111,7 +111,7 @@ class WorkerInterface {
 
   virtual bool IsRegistered() = 0;
 
-  virtual rpc::CoreWorkerClient *rpc_client() = 0;
+  virtual rpc::CoreWorkerClientInterface *rpc_client() = 0;
 };
 
 /// Worker class encapsulates the implementation details of a worker. A worker
@@ -211,7 +211,7 @@ class Worker : public WorkerInterface {
 
   bool IsRegistered() { return rpc_client_ != nullptr; }
 
-  rpc::CoreWorkerClient *rpc_client() {
+  rpc::CoreWorkerClientInterface *rpc_client() {
     RAY_CHECK(IsRegistered());
     return rpc_client_.get();
   }

src/ray/raylet/worker_pool.h

@@ -57,6 +57,28 @@ class WorkerPoolInterface {
   virtual ~WorkerPoolInterface(){};
 };
 
+/// \class IOWorkerPoolInterface
+///
+/// Used for object spilling manager unit tests.
+class IOWorkerPoolInterface {
+ public:
+  /// Add an idle I/O worker to the pool.
+  ///
+  /// \param worker The idle I/O worker to add.
+  virtual void PushIOWorker(const std::shared_ptr<WorkerInterface> &worker) = 0;
+
+  /// Pop an idle I/O worker from the pool and trigger a callback when
+  /// an I/O worker is available.
+  /// The caller is responsible for pushing the worker back onto the
+  /// pool once the worker has completed its work.
+  ///
+  /// \param callback The callback that returns an available I/O worker.
+  virtual void PopIOWorker(
+      std::function<void(std::shared_ptr<WorkerInterface>)> callback) = 0;
+
+  virtual ~IOWorkerPoolInterface(){};
+};
+
 class WorkerInterface;
 class Worker;
 
@@ -64,7 +86,7 @@ class Worker;
 ///
 /// The WorkerPool is responsible for managing a pool of Workers. Each Worker
 /// is a container for a unit of work.
-class WorkerPool : public WorkerPoolInterface {
+class WorkerPool : public WorkerPoolInterface, public IOWorkerPoolInterface {
  public:
   /// Create a pool and asynchronously start at least the specified number of workers per
   /// language.