Add entries to in-memory store on Put() (#7085)

src/ray/core_worker/core_worker.cc

@@ -390,7 +390,7 @@ Status CoreWorker::Put(const RayObject &object,
                        ObjectID *object_id) {
   *object_id = ObjectID::ForPut(worker_context_.GetCurrentTaskID(),
                                 worker_context_.GetNextPutIndex(),
-                                static_cast<uint8_t>(TaskTransportType::RAYLET));
+                                static_cast<uint8_t>(TaskTransportType::DIRECT));
   reference_counter_->AddOwnedObject(*object_id, contained_object_ids, GetCallerId(),
                                      rpc_address_);
   return Put(object, contained_object_ids, *object_id, /*pin_object=*/true);
@@ -419,7 +419,7 @@ Status CoreWorker::Put(const RayObject &object,
       RAY_RETURN_NOT_OK(plasma_store_provider_->Release(object_id));
     }
   }
-  return Status::OK();
+  return memory_store_->Put(RayObject(rpc::ErrorType::OBJECT_IN_PLASMA), object_id);
 }
 
 Status CoreWorker::Create(const std::shared_ptr<Buffer> &metadata, const size_t data_size,
@@ -427,7 +427,7 @@ Status CoreWorker::Create(const std::shared_ptr<Buffer> &metadata, const size_t
                           ObjectID *object_id, std::shared_ptr<Buffer> *data) {
   *object_id = ObjectID::ForPut(worker_context_.GetCurrentTaskID(),
                                 worker_context_.GetNextPutIndex(),
-                                static_cast<uint8_t>(TaskTransportType::RAYLET));
+                                static_cast<uint8_t>(TaskTransportType::DIRECT));
   RAY_RETURN_NOT_OK(
       plasma_store_provider_->Create(metadata, data_size, *object_id, data));
   // Only add the object to the reference counter if it didn't already exist.
@@ -462,7 +462,7 @@ Status CoreWorker::Seal(const ObjectID &object_id, bool pin_object,
   } else {
     RAY_RETURN_NOT_OK(plasma_store_provider_->Release(object_id));
   }
-  return Status::OK();
+  return memory_store_->Put(RayObject(rpc::ErrorType::OBJECT_IN_PLASMA), object_id);
 }
 
 Status CoreWorker::Get(const std::vector<ObjectID> &ids, const int64_t timeout_ms,
@@ -538,13 +538,9 @@ Status CoreWorker::Get(const std::vector<ObjectID> &ids, const int64_t timeout_m
 
 Status CoreWorker::Contains(const ObjectID &object_id, bool *has_object) {
   bool found = false;
-  if (object_id.IsDirectCallType()) {
-    bool in_plasma = false;
-    found = memory_store_->Contains(object_id, &in_plasma);
-    if (in_plasma) {
-      RAY_RETURN_NOT_OK(plasma_store_provider_->Contains(object_id, &found));
-    }
-  } else {
+  bool in_plasma = false;
+  found = memory_store_->Contains(object_id, &in_plasma);
+  if (in_plasma) {
     RAY_RETURN_NOT_OK(plasma_store_provider_->Contains(object_id, &found));
   }
   *has_object = found;
@@ -667,14 +663,9 @@ Status CoreWorker::Delete(const std::vector<ObjectID> &object_ids, bool local_on
 
   // We only delete from plasma, which avoids hangs (issue #7105). In-memory
   // objects are always handled by ref counting only.
-  absl::flat_hash_set<ObjectID> plasma_object_ids;
-  for (const auto &obj_id : object_ids) {
-    plasma_object_ids.insert(obj_id);
-  }
-  RAY_RETURN_NOT_OK(plasma_store_provider_->Delete(plasma_object_ids, local_only,
-                                                   delete_creating_tasks));
-
-  return Status::OK();
+  absl::flat_hash_set<ObjectID> plasma_object_ids(object_ids.begin(), object_ids.end());
+  return plasma_store_provider_->Delete(plasma_object_ids, local_only,
+                                        delete_creating_tasks);
 }
 
 void CoreWorker::TriggerGlobalGC() {
@@ -1407,7 +1398,10 @@ void CoreWorker::HandlePlasmaObjectReady(const rpc::PlasmaObjectReadyRequest &re
                                          rpc::PlasmaObjectReadyReply *reply,
                                          rpc::SendReplyCallback send_reply_callback) {
   RAY_CHECK(plasma_done_callback_ != nullptr) << "Plasma done callback not defined.";
-  // This callback must be asynchronous to allow plasma to receive objects
+  // This callback needs to be asynchronous because it runs on the io_service_, so no
+  // RPCs can be processed while it's running. This can easily lead to deadlock (for
+  // example if the callback calls ray.get() on an object that is dependent on an RPC
+  // to be ready).
   plasma_done_callback_(ObjectID::FromBinary(request.object_id()), request.data_size(),
                         request.metadata_size());
   send_reply_callback(Status::OK(), nullptr, nullptr);

src/ray/core_worker/fiber.h

@@ -129,4 +129,4 @@ class FiberState {
 
 }  // namespace ray
 
-#endif  // RAY_CORE_WORKER_FIBER_H
+#endif  // RAY_CORE_WORKER_FIBER_H

src/ray/core_worker/store_provider/memory_store/memory_store.cc

@@ -158,11 +158,13 @@ std::shared_ptr<RayObject> CoreWorkerMemoryStore::GetOrPromoteToPlasma(
 }
 
 Status CoreWorkerMemoryStore::Put(const RayObject &object, const ObjectID &object_id) {
-  RAY_CHECK(object_id.IsDirectCallType());
   std::vector<std::function<void(std::shared_ptr<RayObject>)>> async_callbacks;
   auto object_entry = std::make_shared<RayObject>(object.GetData(), object.GetMetadata(),
                                                   object.GetNestedIds(), true);
 
+  // TODO(edoakes): we should instead return a flag to the caller to put the object in
+  // plasma.
+  bool should_put_in_plasma = false;
   {
     absl::MutexLock lock(&mu_);
 
@@ -181,11 +183,9 @@ Status CoreWorkerMemoryStore::Put(const RayObject &object, const ObjectID &objec
     auto promoted_it = promoted_to_plasma_.find(object_id);
     if (promoted_it != promoted_to_plasma_.end()) {
       RAY_CHECK(store_in_plasma_ != nullptr);
-      if (!object.IsInPlasmaError()) {
-        // Only need to promote to plasma if it wasn't already put into plasma
-        // by the task that created the object.
-        store_in_plasma_(object, object_id);
-      }
+      // Only need to promote to plasma if it wasn't already put into plasma
+      // by the task that created the object.
+      should_put_in_plasma = !object.IsInPlasmaError();
       promoted_to_plasma_.erase(promoted_it);
     }
 
@@ -212,6 +212,13 @@ Status CoreWorkerMemoryStore::Put(const RayObject &object, const ObjectID &objec
     }
   }
 
+  // Must be called without holding the lock because store_in_plasma_ goes
+  // through the regular CoreWorker::Put() codepath, which calls into the
+  // in-memory store (would cause deadlock).
+  if (should_put_in_plasma) {
+    store_in_plasma_(object, object_id);
+  }
+
   // It's important for performance to run the callbacks outside the lock.
   for (const auto &cb : async_callbacks) {
     cb(object_entry);
@@ -250,7 +257,6 @@ Status CoreWorkerMemoryStore::GetImpl(const std::vector<ObjectID> &object_ids,
       if (iter != objects_.end()) {
         (*results)[i] = iter->second;
         if (remove_after_get) {
-          RAY_LOG(ERROR) << "REMOVE_AFTER_GET";
           // Note that we cannot remove the object_id from `objects_` now,
           // because `object_ids` might have duplicate ids.
           ids_to_remove.insert(object_id);
@@ -437,7 +443,6 @@ bool CoreWorkerMemoryStore::Contains(const ObjectID &object_id, bool *in_plasma)
   if (it != objects_.end()) {
     if (it->second->IsInPlasmaError()) {
       *in_plasma = true;
-      return false;
     }
     return true;
   }

src/ray/core_worker/store_provider/memory_store/memory_store.h

@@ -111,7 +111,7 @@ class CoreWorkerMemoryStore {
   ///
   /// \param[in] object_id The object to check.
   /// \param[out] in_plasma Set to true if the object was spilled to plasma.
-  /// If this is set to true, Contains() will return false.
+  /// Will only be true if the store contains the object.
   /// \return Whether the store has the object.
   bool Contains(const ObjectID &object_id, bool *in_plasma);
 

src/ray/core_worker/test/core_worker_test.cc

@@ -951,76 +951,6 @@ TEST_F(SingleNodeTest, TestObjectInterface) {
   ASSERT_TRUE(!results[1]);
 }
 
-TEST_F(TwoNodeTest, TestObjectInterfaceCrossNodes) {
-  CoreWorker worker1(WorkerType::DRIVER, Language::PYTHON, raylet_store_socket_names_[0],
-                     raylet_socket_names_[0], NextJobId(), gcs_options_, "", "127.0.0.1",
-                     node_manager_port, nullptr);
-
-  CoreWorker worker2(WorkerType::DRIVER, Language::PYTHON, raylet_store_socket_names_[1],
-                     raylet_socket_names_[1], NextJobId(), gcs_options_, "", "127.0.0.1",
-                     node_manager_port, nullptr);
-
-  uint8_t array1[] = {1, 2, 3, 4, 5, 6, 7, 8};
-  uint8_t array2[] = {10, 11, 12, 13, 14, 15};
-
-  std::vector<std::shared_ptr<LocalMemoryBuffer>> buffers;
-  buffers.emplace_back(std::make_shared<LocalMemoryBuffer>(array1, sizeof(array1)));
-  buffers.emplace_back(std::make_shared<LocalMemoryBuffer>(array2, sizeof(array2)));
-
-  std::vector<ObjectID> ids(buffers.size());
-  for (size_t i = 0; i < ids.size(); i++) {
-    RAY_CHECK_OK(worker1.Put(RayObject(buffers[i], nullptr, std::vector<ObjectID>()), {},
-                             &ids[i]));
-  }
-
-  // Test Get() from remote node.
-  std::vector<std::shared_ptr<RayObject>> results;
-  RAY_CHECK_OK(worker2.Get(ids, -1, &results));
-
-  ASSERT_EQ(results.size(), 2);
-  for (size_t i = 0; i < ids.size(); i++) {
-    ASSERT_EQ(results[i]->GetData()->Size(), buffers[i]->Size());
-    ASSERT_EQ(*(results[i]->GetData()), *buffers[i]);
-  }
-
-  // Test Wait() from remote node.
-  ObjectID non_existent_id = ObjectID::FromRandom();
-  std::vector<ObjectID> all_ids(ids);
-  all_ids.push_back(non_existent_id);
-
-  std::vector<bool> wait_results;
-  RAY_CHECK_OK(worker2.Wait(all_ids, 2, -1, &wait_results));
-  ASSERT_EQ(wait_results.size(), 3);
-  ASSERT_EQ(wait_results, std::vector<bool>({true, true, false}));
-
-  RAY_CHECK_OK(worker2.Wait(all_ids, 3, 100, &wait_results));
-  ASSERT_EQ(wait_results.size(), 3);
-  ASSERT_EQ(wait_results, std::vector<bool>({true, true, false}));
-
-  // Test Delete() from all machines.
-  // clear the reference held by PlasmaBuffer.
-  results.clear();
-  RAY_CHECK_OK(worker2.Delete(ids, false, false));
-
-  // Note that Delete() calls RayletClient::FreeObjects and would not
-  // wait for objects being deleted, so wait a while for plasma store
-  // to process the command.
-  usleep(1000 * 1000);
-  // Verify objects are deleted from both machines.
-  ASSERT_TRUE(worker2.Get(ids, 0, &results).IsTimedOut());
-  ASSERT_EQ(results.size(), 2);
-  ASSERT_TRUE(!results[0]);
-  ASSERT_TRUE(!results[1]);
-
-  // TODO(edoakes): this currently fails because the object is pinned on the
-  // creating node. Should be fixed or removed once we decide the semantics
-  // for Delete() with pinning.
-  // ASSERT_TRUE(worker1.Get(ids, 0, &results).IsTimedOut());
-  // ASSERT_EQ(results.size(), 2);
-  // ASSERT_TRUE(!results[0]);
-  // ASSERT_TRUE(!results[1]);
-}
-
 TEST_F(SingleNodeTest, TestNormalTaskLocal) {
   std::unordered_map<std::string, double> resources;
   TestNormalTask(resources);

src/ray/core_worker/test/scheduling_queue_test.cc

@@ -24,7 +24,7 @@ TEST(SchedulingQueueTest, TestInOrder) {
   boost::asio::io_service io_service;
   MockWaiter waiter;
   WorkerContext context(WorkerType::WORKER, JobID::Nil());
-  SchedulingQueue queue(io_service, waiter, context, 0);
+  SchedulingQueue queue(io_service, waiter, context);
   int n_ok = 0;
   int n_rej = 0;
   auto fn_ok = [&n_ok]() { n_ok++; };
@@ -45,7 +45,7 @@ TEST(SchedulingQueueTest, TestWaitForObjects) {
   boost::asio::io_service io_service;
   MockWaiter waiter;
   WorkerContext context(WorkerType::WORKER, JobID::Nil());
-  SchedulingQueue queue(io_service, waiter, context, 0);
+  SchedulingQueue queue(io_service, waiter, context);
   int n_ok = 0;
   int n_rej = 0;
   auto fn_ok = [&n_ok]() { n_ok++; };
@@ -71,7 +71,7 @@ TEST(SchedulingQueueTest, TestWaitForObjectsNotSubjectToSeqTimeout) {
   boost::asio::io_service io_service;
   MockWaiter waiter;
   WorkerContext context(WorkerType::WORKER, JobID::Nil());
-  SchedulingQueue queue(io_service, waiter, context, 0);
+  SchedulingQueue queue(io_service, waiter, context);
   int n_ok = 0;
   int n_rej = 0;
   auto fn_ok = [&n_ok]() { n_ok++; };
@@ -89,7 +89,7 @@ TEST(SchedulingQueueTest, TestOutOfOrder) {
   boost::asio::io_service io_service;
   MockWaiter waiter;
   WorkerContext context(WorkerType::WORKER, JobID::Nil());
-  SchedulingQueue queue(io_service, waiter, context, 0);
+  SchedulingQueue queue(io_service, waiter, context);
   int n_ok = 0;
   int n_rej = 0;
   auto fn_ok = [&n_ok]() { n_ok++; };
@@ -107,7 +107,7 @@ TEST(SchedulingQueueTest, TestSeqWaitTimeout) {
   boost::asio::io_service io_service;
   MockWaiter waiter;
   WorkerContext context(WorkerType::WORKER, JobID::Nil());
-  SchedulingQueue queue(io_service, waiter, context, 0);
+  SchedulingQueue queue(io_service, waiter, context);
   int n_ok = 0;
   int n_rej = 0;
   auto fn_ok = [&n_ok]() { n_ok++; };
@@ -130,7 +130,7 @@ TEST(SchedulingQueueTest, TestSkipAlreadyProcessedByClient) {
   boost::asio::io_service io_service;
   MockWaiter waiter;
   WorkerContext context(WorkerType::WORKER, JobID::Nil());
-  SchedulingQueue queue(io_service, waiter, context, 0);
+  SchedulingQueue queue(io_service, waiter, context);
   int n_ok = 0;
   int n_rej = 0;
   auto fn_ok = [&n_ok]() { n_ok++; };

src/ray/core_worker/transport/direct_actor_transport.cc

@@ -183,7 +183,6 @@ void CoreWorkerDirectTaskReceiver::HandlePushTask(
     rpc::SendReplyCallback send_reply_callback) {
   RAY_CHECK(waiter_ != nullptr) << "Must call init() prior to use";
   const TaskSpecification task_spec(request.task_spec());
-  RAY_LOG(DEBUG) << "Received task " << task_spec.DebugString();
   if (task_spec.IsActorTask() && !worker_context_.CurrentTaskIsDirectCall()) {
     send_reply_callback(Status::Invalid("This actor doesn't accept direct calls."),
                         nullptr, nullptr);

src/ray/raylet/raylet_client.cc

@@ -178,14 +178,6 @@ raylet::RayletClient::RayletClient(
 }
 
 Status raylet::RayletClient::SubmitTask(const TaskSpecification &task_spec) {
-  for (size_t i = 0; i < task_spec.NumArgs(); i++) {
-    if (task_spec.ArgByRef(i)) {
-      for (size_t j = 0; j < task_spec.ArgIdCount(i); j++) {
-        RAY_CHECK(!task_spec.ArgId(i, j).IsDirectCallType())
-            << "Passing direct call objects to non-direct tasks is not allowed.";
-      }
-    }
-  }
   flatbuffers::FlatBufferBuilder fbb;
   auto message =
       protocol::CreateSubmitTaskRequest(fbb, fbb.CreateString(task_spec.Serialize()));