Add ray.internal.free (#2542)

src/local_scheduler/lib/python/local_scheduler_extension.cc

@@ -414,6 +414,43 @@ static PyObject *PyLocalSchedulerClient_push_profile_events(PyObject *self,
   Py_RETURN_NONE;
 }
 
+static PyObject *PyLocalSchedulerClient_free(PyObject *self, PyObject *args) {
+  PyObject *py_object_ids;
+  PyObject *py_local_only;
+
+  if (!PyArg_ParseTuple(args, "OO", &py_object_ids, &py_local_only)) {
+    return NULL;
+  }
+
+  bool local_only = static_cast<bool>(PyObject_IsTrue(py_local_only));
+
+  // Convert object ids.
+  PyObject *iter = PyObject_GetIter(py_object_ids);
+  if (!iter) {
+    return NULL;
+  }
+  std::vector<ObjectID> object_ids;
+  while (true) {
+    PyObject *next = PyIter_Next(iter);
+    ObjectID object_id;
+    if (!next) {
+      break;
+    }
+    if (!PyObjectToUniqueID(next, &object_id)) {
+      // Error parsing object ID.
+      return NULL;
+    }
+    object_ids.push_back(object_id);
+  }
+
+  // Invoke local_scheduler_free_objects_in_object_store.
+  local_scheduler_free_objects_in_object_store(
+      reinterpret_cast<PyLocalSchedulerClient *>(self)
+          ->local_scheduler_connection,
+      object_ids, local_only);
+  Py_RETURN_NONE;
+}
+
 static PyMethodDef PyLocalSchedulerClient_methods[] = {
     {"disconnect", (PyCFunction) PyLocalSchedulerClient_disconnect, METH_NOARGS,
      "Notify the local scheduler that this client is exiting gracefully."},
@@ -446,6 +483,8 @@ static PyMethodDef PyLocalSchedulerClient_methods[] = {
     {"push_profile_events",
      (PyCFunction) PyLocalSchedulerClient_push_profile_events, METH_VARARGS,
      "Store some profiling events in the GCS."},
+    {"free", (PyCFunction) PyLocalSchedulerClient_free, METH_VARARGS,
+     "Free a list of objects from object stores."},
     {NULL} /* Sentinel */
 };
 

src/local_scheduler/local_scheduler_client.cc

@@ -351,3 +351,20 @@ void local_scheduler_push_profile_events(
                     ray::protocol::MessageType::PushProfileEventsRequest),
                 fbb.GetSize(), fbb.GetBufferPointer(), &conn->write_mutex);
 }
+
+void local_scheduler_free_objects_in_object_store(
+    LocalSchedulerConnection *conn,
+    const std::vector<ray::ObjectID> &object_ids,
+    bool local_only) {
+  flatbuffers::FlatBufferBuilder fbb;
+  auto message = ray::protocol::CreateFreeObjectsRequest(
+      fbb, local_only, to_flatbuf(fbb, object_ids));
+  fbb.Finish(message);
+
+  int success = write_message(
+      conn->conn,
+      static_cast<int64_t>(
+          ray::protocol::MessageType::FreeObjectsInObjectStoreRequest),
+      fbb.GetSize(), fbb.GetBufferPointer(), &conn->write_mutex);
+  RAY_CHECK(success == 0) << "Failed to write message to raylet.";
+}

src/local_scheduler/local_scheduler_client.h

@@ -244,4 +244,16 @@ void local_scheduler_push_profile_events(
     LocalSchedulerConnection *conn,
     const ProfileTableDataT &profile_events);
 
+/// Free a list of objects from object stores.
+///
+/// \param conn The connection information.
+/// \param object_ids A list of ObjectsIDs to be deleted.
+/// \param local_only Whether keep this request with local object store
+/// or send it to all the object stores.
+/// \return Void.
+void local_scheduler_free_objects_in_object_store(
+    LocalSchedulerConnection *conn,
+    const std::vector<ray::ObjectID> &object_ids,
+    bool local_only);
+
 #endif

src/ray/gcs/tables.cc

@@ -426,6 +426,10 @@ const ClientTableDataT &ClientTable::GetClient(const ClientID &client_id) const
   }
 }
 
+const std::unordered_map<ClientID, ClientTableDataT> &ClientTable::GetAllClients() const {
+  return client_cache_;
+}
+
 template class Log<ObjectID, ObjectTableData>;
 template class Log<TaskID, ray::protocol::Task>;
 template class Table<TaskID, ray::protocol::Task>;

src/ray/gcs/tables.h

@@ -651,6 +651,13 @@ class ClientTable : private Log<UniqueID, ClientTableData> {
   /// \return Whether the client with ID client_id is removed.
   bool IsRemoved(const ClientID &client_id) const;
 
+  /// Get the information of all clients.
+  ///
+  /// Note: The return value contains ClientID::nil() which should be filtered.
+  ///
+  /// \return The client ID to client information map.
+  const std::unordered_map<ClientID, ClientTableDataT> &GetAllClients() const;
+
  private:
   /// Handle a client table notification.
   void HandleNotification(AsyncGcsClient *client, const ClientTableDataT &notifications);

src/ray/object_manager/format/object_manager.fbs

@@ -4,7 +4,8 @@ namespace ray.object_manager.protocol;
 enum MessageType:int {
   ConnectClient = 1,
   PushRequest,
-  PullRequest
+  PullRequest,
+  FreeRequest
 }
 
 table PushRequestMessage {
@@ -31,3 +32,8 @@ table ConnectClientMessage {
   // Whether this is a transfer connection.
   is_transfer: bool;
 }
+
+table FreeRequestMessage {
+  // List of IDs to be deleted.
+  object_ids: [string];
+}

src/ray/object_manager/object_buffer_pool.cc

@@ -185,4 +185,13 @@ std::vector<ObjectBufferPool::ChunkInfo> ObjectBufferPool::BuildChunks(
   return chunks;
 }
 
+void ObjectBufferPool::FreeObjects(const std::vector<ObjectID> &object_ids) {
+  std::vector<plasma::ObjectID> plasma_ids;
+  plasma_ids.reserve(object_ids.size());
+  for (const auto &id : object_ids) {
+    plasma_ids.push_back(id.to_plasma_id());
+  }
+  ARROW_CHECK_OK(store_client_.Delete(plasma_ids));
+}
+
 }  // namespace ray

src/ray/object_manager/object_buffer_pool.h

@@ -123,6 +123,12 @@ class ObjectBufferPool {
   /// \param chunk_index The index of the chunk.
   void SealChunk(const ObjectID &object_id, uint64_t chunk_index);
 
+  /// Free a list of objects from object store.
+  ///
+  /// \param object_ids the The list of ObjectIDs to be deleted.
+  /// \return Void.
+  void FreeObjects(const std::vector<ObjectID> &object_ids);
+
  private:
   /// Abort the create operation associated with an object. This destroys the buffer
   /// state, including create operations in progress for all chunks of the object.

src/ray/object_manager/object_directory.cc

@@ -117,6 +117,24 @@ ray::Status ObjectDirectory::GetInformation(const ClientID &client_id,
   return ray::Status::OK();
 }
 
+void ObjectDirectory::RunFunctionForEachClient(
+    const InfoSuccessCallback &client_function) {
+  const auto &clients = gcs_client_->client_table().GetAllClients();
+  for (const auto &client_pair : clients) {
+    const ClientTableDataT &data = client_pair.second;
+    if (client_pair.first == ClientID::nil() ||
+        client_pair.first == gcs_client_->client_table().GetLocalClientId() ||
+        !data.is_insertion) {
+      continue;
+    } else {
+      const auto &info =
+          RemoteConnectionInfo(client_pair.first, data.node_manager_address,
+                               static_cast<uint16_t>(data.object_manager_port));
+      client_function(info);
+    }
+  }
+}
+
 ray::Status ObjectDirectory::SubscribeObjectLocations(const UniqueID &callback_id,
                                                       const ObjectID &object_id,
                                                       const OnLocationsFound &callback) {

src/ray/object_manager/object_directory.h

@@ -101,6 +101,13 @@ class ObjectDirectoryInterface {
   /// \return Status of whether this method succeeded.
   virtual ray::Status ReportObjectRemoved(const ObjectID &object_id,
                                           const ClientID &client_id) = 0;
+
+  /// Go through all the client information.
+  ///
+  /// \param success_cb A callback which handles the success of this method.
+  /// This function will be called multiple times.
+  /// \return Void.
+  virtual void RunFunctionForEachClient(const InfoSuccessCallback &client_function) = 0;
 };
 
 /// Ray ObjectDirectory declaration.
@@ -115,6 +122,8 @@ class ObjectDirectory : public ObjectDirectoryInterface {
                              const InfoSuccessCallback &success_callback,
                              const InfoFailureCallback &fail_callback) override;
 
+  void RunFunctionForEachClient(const InfoSuccessCallback &client_function) override;
+
   ray::Status LookupLocations(const ObjectID &object_id,
                               const OnLocationsFound &callback) override;
 

src/ray/object_manager/object_manager.cc

@@ -1,4 +1,5 @@
 #include "ray/object_manager/object_manager.h"
+#include "common/common_protocol.h"
 #include "ray/util/util.h"
 
 namespace asio = boost::asio;
@@ -655,6 +656,10 @@ void ObjectManager::ProcessClientMessage(std::shared_ptr<TcpClientConnection> &c
     ConnectClient(conn, message);
     break;
   }
+  case static_cast<int64_t>(object_manager_protocol::MessageType::FreeRequest): {
+    ReceiveFreeRequest(conn, message);
+    break;
+  }
   case static_cast<int64_t>(protocol::MessageType::DisconnectClient): {
     // TODO(hme): Disconnect without depending on the node manager protocol.
     DisconnectClient(conn, message);
@@ -755,4 +760,51 @@ void ObjectManager::ExecuteReceiveObject(const ClientID &client_id,
                  << "/" << config_.max_receives;
 }
 
+void ObjectManager::ReceiveFreeRequest(std::shared_ptr<TcpClientConnection> &conn,
+                                       const uint8_t *message) {
+  auto free_request =
+      flatbuffers::GetRoot<object_manager_protocol::FreeRequestMessage>(message);
+  std::vector<ObjectID> object_ids = from_flatbuf(*free_request->object_ids());
+  // This RPC should come from another Object Manager.
+  // Keep this request local.
+  bool local_only = true;
+  FreeObjects(object_ids, local_only);
+  conn->ProcessMessages();
+}
+
+void ObjectManager::FreeObjects(const std::vector<ObjectID> &object_ids,
+                                bool local_only) {
+  buffer_pool_.FreeObjects(object_ids);
+  if (!local_only) {
+    SpreadFreeObjectRequest(object_ids);
+  }
+}
+
+void ObjectManager::SpreadFreeObjectRequest(const std::vector<ObjectID> &object_ids) {
+  // This code path should be called from node manager.
+  flatbuffers::FlatBufferBuilder fbb;
+  flatbuffers::Offset<object_manager_protocol::FreeRequestMessage> request =
+      object_manager_protocol::CreateFreeRequestMessage(fbb, to_flatbuf(fbb, object_ids));
+  fbb.Finish(request);
+  auto function_on_client = [this, &fbb](const RemoteConnectionInfo &connection_info) {
+    std::shared_ptr<SenderConnection> conn;
+    connection_pool_.GetSender(ConnectionPool::ConnectionType::MESSAGE,
+                               connection_info.client_id, &conn);
+    if (conn == nullptr) {
+      conn = CreateSenderConnection(ConnectionPool::ConnectionType::MESSAGE,
+                                    connection_info);
+      connection_pool_.RegisterSender(ConnectionPool::ConnectionType::MESSAGE,
+                                      connection_info.client_id, conn);
+    }
+    ray::Status status = conn->WriteMessage(
+        static_cast<int64_t>(object_manager_protocol::MessageType::FreeRequest),
+        fbb.GetSize(), fbb.GetBufferPointer());
+    if (status.ok()) {
+      connection_pool_.ReleaseSender(ConnectionPool::ConnectionType::MESSAGE, conn);
+    }
+    // TODO(Yuhong): Implement ConnectionPool::RemoveSender and call it in "else".
+  };
+  object_directory_->RunFunctionForEachClient(function_on_client);
+}
+
 }  // namespace ray

src/ray/object_manager/object_manager.h

@@ -163,6 +163,13 @@ class ObjectManager : public ObjectManagerInterface {
                    uint64_t num_required_objects, bool wait_local,
                    const WaitCallback &callback);
 
+  /// Free a list of objects from object store.
+  ///
+  /// \param object_ids the The list of ObjectIDs to be deleted.
+  /// \param local_only Whether keep this request with local object store
+  ///                   or send it to all the object stores.
+  void FreeObjects(const std::vector<ObjectID> &object_ids, bool local_only);
+
  private:
   friend class TestObjectManager;
 
@@ -214,6 +221,11 @@ class ObjectManager : public ObjectManagerInterface {
   /// Completion handler for Wait.
   void WaitComplete(const UniqueID &wait_id);
 
+  /// Spread the Free request to all objects managers.
+  ///
+  /// \param object_ids the The list of ObjectIDs to be deleted.
+  void SpreadFreeObjectRequest(const std::vector<ObjectID> &object_ids);
+
   /// Handle starting, running, and stopping asio io_service.
   void StartIOService();
   void RunSendService();
@@ -271,6 +283,9 @@ class ObjectManager : public ObjectManagerInterface {
   /// Handles receiving a pull request message.
   void ReceivePullRequest(std::shared_ptr<TcpClientConnection> &conn,
                           const uint8_t *message);
+  /// Handles freeing objects request.
+  void ReceiveFreeRequest(std::shared_ptr<TcpClientConnection> &conn,
+                          const uint8_t *message);
 
   /// Handles connect message of a new client connection.
   void ConnectClient(std::shared_ptr<TcpClientConnection> &conn, const uint8_t *message);

src/ray/raylet/format/node_manager.fbs

@@ -68,6 +68,8 @@ enum MessageType:int {
   // Push some profiling events to the GCS. When sending this message to the
   // node manager, the message itself is serialized as a ProfileTableData object.
   PushProfileEventsRequest,
+  // Free the objects in objects store.
+  FreeObjectsInObjectStoreRequest,
 }
 
 table TaskExecutionSpecification {
@@ -177,3 +179,11 @@ table PushErrorRequest {
   // The timestamp of the error message.
   timestamp: double;
 }
+
+table FreeObjectsRequest {
+  // Whether keep this request with local object store 
+  // or send it to all the object stores.
+  local_only: bool;
+  // List of object ids we'll delete from object store.
+  object_ids: [string];
+}

src/ray/raylet/node_manager.cc

@@ -712,6 +712,11 @@ void NodeManager::ProcessClientMessage(
 
     RAY_CHECK_OK(gcs_client_->profile_table().AddProfileEventBatch(*message));
   } break;
+  case protocol::MessageType::FreeObjectsInObjectStoreRequest: {
+    auto message = flatbuffers::GetRoot<protocol::FreeObjectsRequest>(message_data);
+    std::vector<ObjectID> object_ids = from_flatbuf(*message->object_ids());
+    object_manager_.FreeObjects(object_ids, message->local_only());
+  } break;
 
   default:
     RAY_LOG(FATAL) << "Received unexpected message type " << message_type;

src/ray/raylet/reconstruction_policy_test.cc

@@ -47,6 +47,7 @@ class MockObjectDirectory : public ObjectDirectoryInterface {
   MOCK_METHOD3(ReportObjectAdded,
                ray::Status(const ObjectID &, const ClientID &, const ObjectInfoT &));
   MOCK_METHOD2(ReportObjectRemoved, ray::Status(const ObjectID &, const ClientID &));
+  MOCK_METHOD1(RunFunctionForEachClient, void(const InfoSuccessCallback &success_cb));
 
  private:
   std::vector<std::pair<ObjectID, OnLocationsFound>> callbacks_;